#!/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

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: None | 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 fps != fps2:
    #     raise ValueError("FPS of both Videos must be the same (May not work yet, more testing needs to be done)")
    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, offset = find_offset(Quality_Video, German_Video)
        print(f"Offset: {n1} frames| {offset}s")

    if offset > duration or offset > duration2:
        raise ValueError("Offset needs to be smaller than the video durations")

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

    # --- 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 offset >= 0:
    #     offset = -offset
    #     # Quality starts later than German
    #     start = vstart + offset
    #     if duration < offset + duration2:
    #         # When the german audio is shorter than the original video
    #         end = start + duration
    #         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
    #         start2 = vstart2 + (end - vstart)
    #         end2 = vstart2 + duration2
    #         if end2 > start2:
    #             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)
    #     else:
    #         end = 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)
    # else:
        # German starts Later than Quality
        # Get the start from the original Audio
    start2 = vstart2
    end2 = start2 + offset
    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)
    ###
    # if duration - offset < duration2:
        # When the german audio is shorter than the original video
    start = vstart
    end = start + duration
    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
    start2 = vstart2 + end
    # end2 = start2 + (duration2 - duration - (n1 / fps))
    end2 = vstart2 + get_audio_duration(Quality_Video)
    # print(f"duration {duration}")
    # print(f"duration2 {duration2}")
    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)
    # else:
    #     start = vstart
    #     end = 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)

    # german_audio_file2 = output_dir / f"{Path(Quality_Video).stem}_german_audio2.mka"
    # german_audio_file3 = output_dir / f"{Path(Quality_Video).stem}_german_audio3.mka"

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

    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 = "/home/honney/mount/Media/tmp/serien/ready/Maid-sama (2010) [tmdb-65648]/Season 01"
    start_num = 17
    episode_num = 26
    for i in range(start_num, episode_num+1):
        # --- German_Videouts ---
        pattern = f"/home/honney/mount/Media/tmp/serien/Maid Sama All/Maid Sama/S01/{i:02}.mkv"
        print(f"German_Video={pattern}")
        German_Video = glob.glob(pattern)
        pattern = f"/home/honney/mount/Media/tmp/serien/Maid Sama All/Kaichou wa Maid-sama!/S01/{i:02}.mkv"
        Quality_Video = glob.glob(pattern)
        print(f"German_Video={German_Video}")
        print(f"Quality_Video={Quality_Video}")

        if len(German_Video) > 1:
            raise ValueError("There are to many videos matching the German wildcard")
        if len(Quality_Video) > 1:
            raise ValueError("There are to many videos matching the Quality wildcard")
        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()