#!/usr/bin/env python3

import subprocess
from pathlib import Path
import cv2
import numpy as np
from scipy import signal

def run(cmd):
    """Run a subprocess command and raise on error."""
    subprocess.run(cmd, check=True)

def get_video_signal(path):
    cap = cv2.VideoCapture(path)
    motion_signal = []
    prev_gray = None
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        # 1. Shrink the frame (makes it faster and removes noise)
        small_frame = cv2.resize(frame, (160, 90))
        gray = cv2.cvtColor(small_frame, cv2.COLOR_BGR2GRAY)
        
        if prev_gray is not None:
            # 2. Calculate the difference between this frame and the last
            # This detects the "movement" or "cuts"
            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)

def find_offset(video_a, video_b):
    # print("Extracting signals (this may take a moment)...")
    sig_a = get_video_signal(video_a)
    # print("sig_a finished")
    sig_b = get_video_signal(video_b)
    # print("sig_b finished")

    # Normalize signals (subtract mean, divide by std dev) to improve correlation
    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)

    # Perform Cross-Correlation
    correlation = signal.correlate(sig_a, sig_b, mode='full')
    lags = signal.correlation_lags(len(sig_a), len(sig_b), mode='full')
    # plt.plot(lags)
    
    # The peak of the correlation is our offset in frames
    best_lag = lags[np.argmax(correlation)]
    
    return best_lag

def probe_video(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 trim_audio(input_mkv, start_time, end_time, out_audio):
    """Trim only the first (and only) audio track from input_mkv."""
    run([
        "ffmpeg",
        "-y",
        "-i", input_mkv,
        "-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(inp2, audio_file, output_dir):
    """Mux trimmed audio into inp2 and place final file in output_dir."""
    output_dir = Path(output_dir)
    output_dir.mkdir(exist_ok=True)

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

    run([
        "mkvmerge",
        "-o", str(final_output),
        str(inp2),
        "--language", "0:deu",  # assign German language
        "--track-name", "0:German (2.0 aac)",
        str(audio_file)
    ])

    return final_output

def main():
    for i in range(1, 12):
        # --- Inputs ---
        inp = 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"
        inp2 = f"/home/honney/projects/mkv/Anohana - The Flower We Saw That Day/S01/AnoHana The Flower We Saw That Day S01E0{i:02}.mkv"
        n1 = find_offset(inp, inp2)                  # first frame to keep
        output_dir = "./output_dir"

        # --- Probe video ---
        fps, _, _, vstart = probe_video(inp)
        _, duration2, _, _ = probe_video(inp2)

        # --- Compute timestamps ---
        start = vstart + (n1 / fps)
        end = start + duration2

        if end <= start:
            raise ValueError("n2 must be greater than n1")

        print(f"fps: {fps}")
        print(f"Cutting audio from frame {n1} (timestamp {start:.6f}s) to timestamp {end:.6f}s")

        # --- Trim only audio ---
        audio_file = Path(output_dir) / f"{Path(inp2).stem}_german_audio.aac"
        trim_audio(inp, start, end, audio_file)

        # --- Mux trimmed audio ---
        final_file = mux_audio_into_output(inp2, audio_file, output_dir)

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

if __name__ == "__main__":
    main()