#!/usr/bin/env python3

import subprocess
from pathlib import Path
import cv2
import numpy as np
from scipy import signal
import sys
import os
import glob
import tempfile
import shutil
from typing import Tuple, Optional

DEBUG = True
TARGET_FPS = 30

def run(cmd):
    """
    Run a subprocess command.

    Args:
        cmd: Command and arguments as a list.

    Raises:
        subprocess.CalledProcessError: If the command exits with a non-zero status.
    """
    try:
        result = subprocess.run(
            cmd,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True
        )
        return result
    except subprocess.CalledProcessError as e:
        print(e.stdout)
        raise

def get_video_signal(path: Path) -> Tuple[np.ndarray, float]:
    cap = cv2.VideoCapture(path)
    fps = cap.get(cv2.CAP_PROP_FPS)

    motion_signal = []
    prev_gray = None

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        small_frame = cv2.resize(frame, (160, 90))
        gray = cv2.cvtColor(small_frame, cv2.COLOR_BGR2GRAY)

        if prev_gray is not None:
            diff = cv2.absdiff(gray, prev_gray)
            motion_signal.append(np.sum(diff))
        else:
            motion_signal.append(0)

        prev_gray = gray

    cap.release()

    return np.array(motion_signal), fps


def resample_signal(sig: np.ndarray, original_fps: float, target_fps: float):
    duration = len(sig) / original_fps
    target_length = int(duration * target_fps)

    return signal.resample(sig, target_length)


def find_offset(video_a: Path, video_b: Path) -> Tuple[int, float]:

    sig_a, fps_a = get_video_signal(video_a)
    sig_b, fps_b = get_video_signal(video_b)

    # Resample both signals to common FPS
    sig_a = resample_signal(sig_a, fps_a, TARGET_FPS)
    sig_b = resample_signal(sig_b, fps_b, TARGET_FPS)

    # Normalize
    sig_a = (sig_a - np.mean(sig_a)) / (np.std(sig_a) + 1e-5)
    sig_b = (sig_b - np.mean(sig_b)) / (np.std(sig_b) + 1e-5)

    correlation = signal.correlate(sig_a, sig_b, mode='full')
    lags = signal.correlation_lags(len(sig_a), len(sig_b), mode='full')

    best_lag = lags[np.argmax(correlation)]

    # Convert lag to seconds
    offset_seconds = best_lag / TARGET_FPS

    return best_lag, offset_seconds

def probe_video(path: Path):
    """Get video FPS, duration, codec, and stream start time."""
    import json
    cmd = [
        "ffprobe",
        "-v", "error",
        "-print_format", "json",
        "-show_streams",
        "-show_format",
        path
    ]
    data = json.loads(subprocess.run(cmd, capture_output=True, text=True, check=True).stdout)

    duration = float(data["format"]["duration"])
    for s in data["streams"]:
        if s["codec_type"] == "video":
            num, den = s["r_frame_rate"].split("/")
            fps = float(num) / float(den)
            codec = s["codec_name"]
            start_time = float(s.get("start_time", 0))
            return fps, duration, codec, start_time

    raise RuntimeError("No video stream found")

def get_audio_duration(file_path: Path):
    """
    Get the actual duration of the first audio stream by asking ffprobe
    to count packets/frames if the stream metadata is missing.
    Returns duration in seconds.
    """

    # As a last resort, use the format duration
    result = subprocess.run([
        "ffprobe",
        "-v", "error",
        "-select_streams", "a:0",
        "-show_entries", "stream_tags=DURATION",
        "-of", "default=noprint_wrappers=1:nokey=1",
        file_path
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )
    result = result.stdout.strip().split(":")[::-1]
    # if DEBUG: print(f"audio duration: {result}")
    res = 0
    for i in range(len(result)):
        res += float(result[i])*(60**i)
    return res

def trim_audio(video: Path, start_time: float, end_time: float, out_audio: Path):
    """
    Extract a segment of the first audio track from a video file.

    Uses FFmpeg to copy the selected audio range without re-encoding,
    which makes the operation fast and preserves the original audio quality.

    Args:
        German_Videout_mkv: Path to the source video file.
        start_time: Start time of the audio segment (seconds or FFmpeg time format).
        end_time: End time of the audio segment.
        out_audio: Path to the output audio file.

    Notes:
        The function selects the first audio stream (`0:a:0`) and copies it
        directly (`-c:a copy`), meaning the codec and quality remain unchanged.
    """
    run([
        "ffmpeg",
        "-y",
        "-i", video,
        "-ss", str(start_time),
        "-to", str(end_time),
        "-map", "0:a:0",      # first audio track
        "-c:a", "copy",       # no re-encode
        out_audio
    ])

def mux_audio_into_output(Quality_Video, audio_file, output_dir):
    """Mux trimmed audio into Quality_Video and place final file in output_dir."""
    output_dir = Path(output_dir)
    output_dir.mkdir(exist_ok=True)

    # final_output = output_dir / f"{Path(Quality_Video).stem}_ger.mkv"
    final_output = output_dir / f"{Path(Quality_Video).stem}.mkv"

    run([
        "mkvmerge",
        "-o", str(final_output),
        str(Quality_Video),
        "--language", "0:deu",  # assign German language
        "--track-name", "0:German (Opus Stereo)",
        str(audio_file)
    ])

    return final_output

def stick_together_audio(left_audio: Path, right_audio: Path, output_audio: Path) -> None:
    """
    Concatenate two audio files so that left_audio plays directly
    before right_audio in the output file.
    """

    left_audio = Path(left_audio)
    right_audio = Path(right_audio)
    output_audio = Path(output_audio)

    # Use stem instead of prefix to create a temporary file
    output_audio_tmp = output_audio.parent / (output_audio.stem + "_temp" + output_audio.suffix)

    run([
        "ffmpeg",
        "-y",
        "-i", str(left_audio),
        "-i", str(right_audio),
        "-filter_complex", "[0:a][1:a]concat=n=2:v=0:a=1[a]",
        "-map", "[a]",
        "-c:a", "opus",   # or "aac" if you prefer
        "-strict", "-2",
        str(output_audio_tmp)
    ])

    shutil.move(output_audio_tmp, output_audio)

def process_one(German_Video: Path, Quality_Video: Path, output_dir: Path, offset: Optional[int] = None):

    print(f"Working on {Quality_Video}")

    # --- Probe video ---
    fps, duration, _, vstart = probe_video(German_Video)
    fps2, duration2, _, vstart2 = probe_video(Quality_Video)

    if duration <= 0:
        raise ValueError("Duration of German_Video must be greater than 0")
    if duration2 <= 0:
        raise ValueError("Duration of Quality_Video must be greater than 0")
    if offset == None:
        n1, offset0 = find_offset(Quality_Video, German_Video)
        print(f"Offset: {n1} frames| {offset0}s")
    else:
        offset0 = offset
    if offset0 > duration or offset0 > duration2:
        raise ValueError("Offset needs to be smaller than the video durations")

    # --- Trim only audio ---
    german_audio_file = output_dir / f"{Path(Quality_Video).stem}_german_audio.mka"
    quality_audio_file1 = output_dir / f"{Path(Quality_Video).stem}_original_audio1.mka"
    quality_audio_file2 = output_dir / f"{Path(Quality_Video).stem}_original_audio2.mka"

    if os.path.isfile(quality_audio_file1):
        os.remove(quality_audio_file1)
    if os.path.isfile(quality_audio_file2):
        os.remove(quality_audio_file2)

    ### Get the first part till german starts (if german does not start before file)
    start2 = vstart2
    end2 = start2 + offset0
    if end2 - start2 > 0.05:
        if DEBUG: print(f"Cutting audio ({quality_audio_file1}) (from timestamp {start2:.6f}s) to timestamp {end2:.6f}s")
        trim_audio(Quality_Video, start2, end2, quality_audio_file1)

    ### Get the German Audio
    ### TODO: get the offset off if german starts before the original
    if offset0 < 0:
        start = vstart - offset0
    else:
        start = vstart
    end = min(vstart + duration, start + duration2)
    if DEBUG: print(f"Cutting audio ({german_audio_file}) (from timestamp {start:.6f}s) to timestamp {end:.6f}s")
    trim_audio(German_Video, start, end, german_audio_file)

    ### Get the end from the first original audio if it is longer than the german part
    start2 = vstart2 + end
    # end2 = start2 + (duration2 - duration - (n1 / fps))
    end2 = vstart2 + get_audio_duration(Quality_Video)
    if end2 - start2 > 0.5:
        if DEBUG: print(f"Cutting audio ({quality_audio_file2}) (from timestamp {start2:.6f}s) to timestamp {end2:.6f}s")
        trim_audio(Quality_Video, start2, end2, quality_audio_file2)

    if os.path.isfile(quality_audio_file1):
        stick_together_audio(quality_audio_file1, german_audio_file, german_audio_file)
    if os.path.isfile(quality_audio_file2):
        stick_together_audio(german_audio_file, quality_audio_file2, german_audio_file)

    # --- Mux trimmed audio ---
    final_file = mux_audio_into_output(Quality_Video, german_audio_file, output_dir)

    print(f"Final file with trimmed audio created at: {final_file}")

    print(f"Removing tmp Audiofiles")
    if os.path.isfile(german_audio_file):
        os.remove(german_audio_file)
    if os.path.isfile(quality_audio_file1):
        os.remove(quality_audio_file1)
    if os.path.isfile(quality_audio_file2):
        os.remove(quality_audio_file2)

def main():
    output_dir = "./output_dir"
    start_num = 1
    episode_num = 11
    for i in range(start_num, episode_num+1):
    # for i in [7, 8, 10, 11, 13, 16, 17, 19, 20, 22, 25]:
    # for i in [7]:
        # --- German_Videouts ---
        german_pattern = f"/home/honney/projects/mkv/AnoHana The Flower We Saw That Day (2011-2011) [imdbid-tt1913273]/Season 01/AnoHana The Flower We Saw That Day S01E0{i:02}.mkv"
        if os.path.isfile(german_pattern):
            German_Video = [german_pattern]
        else: 
            German_Video = glob.glob(german_pattern)
        quality_pattern = f"/home/honney/projects/mkv/Anohana - The Flower We Saw That Day/S01/AnoHana The Flower We Saw That Day S01E0{i:02}.mkv"
        if os.path.isfile(quality_pattern):
            Quality_Video = [quality_pattern]
        else: 
            Quality_Video = glob.glob(quality_pattern)
        print(f"German_Video={German_Video}")
        print(f"Quality_Video={Quality_Video}")

        if len(German_Video) != 1:
            raise ValueError(f"Error with German wildcard:\n{german_pattern}\n{German_Video}")
        if len(Quality_Video) != 1:
            raise ValueError(f"Error with Quality wildcard:\n{quality_pattern}\n{Quality_Video}")
        process_one(Path(German_Video[0]), Path(Quality_Video[0]), Path(output_dir))
        # sys.exit()

if __name__ == "__main__":
    if len(sys.argv) == 3:
        print(sys.argv)
        print(f"German_Video={sys.argv[1]}")
        print(f"Quality_Video={sys.argv[2]}")
        process_one(German_Video=Path(sys.argv[1]), Quality_Video=Path(sys.argv[2]), output_dir=Path("./output_dir"))
    elif len(sys.argv) == 4:
        print(sys.argv)
        print(f"German_Video={sys.argv[1]}")
        print(f"Quality_Video={sys.argv[2]}")
        process_one(German_Video=Path(sys.argv[1]), Quality_Video=Path(sys.argv[2]), output_dir=Path(sys.argv[3]))
    else:
        # print(len(sys.argv), sys.argv)
        # sys.exit()
        main()