#!/bin/bash

# Function to convert timestamp to seconds, including milliseconds
timestamp_to_seconds() {
    IFS=':' read -r hours minutes seconds <<< "$1"
    IFS='.' read -r seconds milliseconds <<< "$seconds"
    total_seconds=$(bc <<< "$hours * 3600 + $minutes * 60 + $seconds")
    total_milliseconds=$(bc <<< "$milliseconds / 1000")
    echo "$total_seconds + $total_milliseconds" | bc
}

# Function to convert seconds back to HH:MM:SS.mmm format for mkvmerge
seconds_to_hms() {
    total_seconds="$1"
    # Separate whole seconds and milliseconds
    seconds=$(echo "$total_seconds" | cut -d'.' -f1)
    milliseconds=$(echo "$total_seconds" | cut -d'.' -f2 | sed 's/^[0-9]\{1,3\}$/&000/')  # Add trailing zeros if necessary
    
    # Ensure milliseconds have 3 digits
    milliseconds=$(echo "$milliseconds" | cut -c1-3)

    # Convert seconds to HH:MM:SS format
    hours=$(bc <<< "$seconds / 3600")
    minutes=$(bc <<< "($seconds % 3600) / 60")
    seconds=$(bc <<< "$seconds % 60")

    # Print in HH:MM:SS.mmm format
    printf "%02d:%02d:%02d.%03d" "$hours" "$minutes" "$seconds" "$milliseconds"
}

# Function to split the video using mkvmerge
split_video() {
    input_file="$1"
    start_time="$2"
    end_time="$3"
    output_file="$4"
    
    start_hms=$(seconds_to_hms "$start_time")
    end_hms=$(seconds_to_hms "$end_time")
    
    # Print the mkvmerge command before running it
    echo "Running command: mkvmerge -o \"$output_file\" --split parts:\"$start_hms\"-\"$end_hms\" \"$input_file\""

    # Execute the command
    mkvmerge -o "$output_file" --split parts:"$start_hms"-"$end_hms" "$input_file"
}

# Function to get video duration
get_duration() {
    input_file="$1"
    duration=$(ffprobe -v quiet -print_format json -show_format "$input_file" | jq -r '.format.duration')
    echo "$duration"
}

# Function to extract chapter start and end times from the input file
get_chapters() {
    input_file="$1"
    chapters=()
    
    # Extract the chapters using mkvinfo, then parse the start and end times
    while IFS= read -r line; do
        if [[ "$line" =~ "start" ]]; then
            start_time=$(echo "$line" | awk '{print $2}')
            chapters+=("$start_time")
        fi
    done < <(mkvinfo "$input_file" | grep "Chapter")

    echo "${chapters[@]}"
}

# Main function
main() {
    only_between=false
    
    if [ "$1" == "--only_between" ]; then
        only_between=true
        shift
    fi
    
    if [ "$#" -lt 3 ]; then
        echo "Usage: $0 [--only_between] <input_file> <timestamp1> <timestamp2> ..."
        exit 1
    fi

    input_file="$1"
    shift
    timestamps=($@)
    output_folder="split_output"
    mkdir -p "$output_folder"

    echo "Splitting video from provided timestamps..."
    
    # Handle chapter timecodes if present
    if [[ "${timestamps[0]}" == "chapters" ]]; then
        # Extract chapter start times
        chapter_times=($(get_chapters "$input_file"))
        
        echo "Extracting chapters from the input file..."
        
        for ((i = 0; i < ${#chapter_times[@]} - 1; i++)); do
            start_time="${chapter_times[$i]}"
            end_time="${chapter_times[$i+1]}"
            output_filename="$output_folder/chapter_$((i+1)).mkv"
            echo "Splitting from Chapter $((i+1)) time $start_time to Chapter $((i+2)) time $end_time..."
            split_video "$input_file" "$start_time" "$end_time" "$output_filename"
        done
    elif $only_between; then
        echo "Extracting only specified segments..."
        for ((i = 0; i < ${#timestamps[@]} - 1; i++)); do
            start_time=$(timestamp_to_seconds "${timestamps[$i]}")
            end_time=$(timestamp_to_seconds "${timestamps[$i+1]}")
            output_filename="$output_folder/part_$((i+1)).mkv"
            echo "Splitting from ${timestamps[$i]} to ${timestamps[$i+1]}..."
            split_video "$input_file" "$start_time" "$end_time" "$output_filename"
        done
    else
        echo "Extracting all segments..."
        start_time=0
        output_index=1
        for timestamp in "${timestamps[@]}"; do
            end_time=$(timestamp_to_seconds "$timestamp")
            output_filename="$output_folder/part_${output_index}.mkv"
            echo "Splitting from $start_time to $timestamp..."
            split_video "$input_file" "$start_time" "$end_time" "$output_filename"
            start_time=$end_time
            ((output_index++))
        done
        duration=$(get_duration "$input_file")
        output_filename="$output_folder/part_${output_index}.mkv"
        echo "Splitting from $start_time to $duration..."
        split_video "$input_file" "$start_time" "$duration" "$output_filename"
    fi

    echo "Splitting complete. Files saved in $output_folder."
}

main "$@"