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

DEBUG = True

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):
    """
    Compute a simple motion signal from a video.

    The function reads a video frame-by-frame and measures the amount of
    pixel change between consecutive frames. Each frame is resized and
    converted to grayscale before calculating the absolute difference
    from the previous frame. The summed pixel differences form a 1-D
    motion signal representing how much visual change occurs over time.

    Args:
        path: Path to the video file to analyze.

    Returns:
        A NumPy array where each element represents the magnitude of
        frame-to-frame change for a given frame. Higher values indicate
        stronger motion or scene changes.

    Raises:
        RuntimeError: If the video cannot be opened or read.

    Notes:
        Frames are resized to 160×90 before processing to reduce noise
        and computational cost.
    """
    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: Path, video_b: Path) -> int:
    """
    Estimate the temporal offset between two videos.

    The function extracts a motion signal from each video using
    `get_video_signal` and then computes the cross-correlation between
    the two signals. The lag corresponding to the maximum correlation
    indicates the frame offset where the videos align best.

    Args:
        video_a: Path to the first video file.
        video_b: Path to the second video file.

    Returns:
        The offset in frames between the two videos. A positive value
        means `video_b` starts later than `video_a`, while a negative
        value means `video_b` starts earlier.

    Raises:
        RuntimeError: If either input path does not point to a valid file.

    Notes:
        Before computing the correlation, both signals are normalized
        (mean subtracted and divided by standard deviation) to make the
        comparison robust to differences in overall motion intensity.
    """
    if not os.path.isfile(video_a):
        raise RuntimeError("video_a is not a file")
    if not os.path.isfile(video_b):
        raise RuntimeError("video_b is not a file")

    # 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: 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(German_Videout_mkv: 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", German_Videout_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(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 (2.0 aac)",
        str(audio_file)
    ])

    return final_output

def stick_together_audio(left_audio: Path, right_audio: Path, output_audio: Path) -> None:
    """
    Concatenate two audio files end-to-end into a single file.

    Args:
        left_audio: Path to the first audio file.
        right_audio: Path to the second audio file.
        output_audio: Path to the resulting audio file. Can overwrite
                      one of the input files safely.

    Notes:
        Works for audio files with the same codec and format. Internally,
        writes to a temporary file first, then moves it to `output_audio`
        to avoid overwriting input files prematurely.
    """
    left_audio = Path(left_audio)
    right_audio = Path(right_audio)
    output_audio = Path(output_audio)

    # Temporary output file in the same directory to allow atomic move
    tmp_output = output_audio.with_name(f"{output_audio.stem}_tmp{output_audio.suffix}")

    # Create a temporary file list for FFmpeg
    with tempfile.NamedTemporaryFile("w", delete=False, suffix=".txt") as f:
        f.write(f"file '{left_audio.resolve()}'\n")
        f.write(f"file '{right_audio.resolve()}'\n")
        list_file = Path(f.name)

    try:
        run([
            "ffmpeg",
            "-y",                # overwrite temporary output if exists
            "-f", "concat",      # concat demuxer
            "-safe", "0",        # allow absolute paths
            "-i", str(list_file),
            "-c", "copy",        # copy streams without re-encoding
            str(tmp_output)
        ])

        # Move temporary output to the final location
        shutil.move(str(tmp_output), str(output_audio))
    finally:
        # Clean up the temporary file list
        if list_file.exists():
            list_file.unlink()
        # Remove temporary output if something went wrong
        if tmp_output.exists() and not output_audio.exists():
            tmp_output.unlink()

def process_one(German_Video: Path, Quality_Video: Path, output_dir: Path):
    # --- 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")
    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")
    
    # n1 = find_offset(German_Video, Quality_Video)
    n1 = -1
    print(f"Offset: {n1}")

    # 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 n1 < 0:
        n1 = -n1
        # Quality starts later than German
        start = vstart + (n1 / fps)
        if duration < (n1 / fps) + 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 + (n1 / fps)
        if end2 > start2:
            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 - (n1 / fps) < 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))
            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:
            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)

    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/Black Bullet/Season 01"
    for i in range(1, 14):
        # --- German_Videouts ---
        # German_Video = f"/home/honney/mount/Media/tmp/serien/ger_low_quality/Black Bullet/Black Bullet Season 01 E{i:02}*.mp4"
        # Quality_Video = f"/home/honney/mount/Media/tmp/serien/high_quality/Black Bullet (Season 1)/Season 01/Black Bullet - S01E{i:02}.mkv"
        pattern = f"/home/honney/mount/Media/tmp/serien/ger_low_quality/Black Bullet/Black Bullet Season 01 E{i:02}*.mp4"
        German_Video = glob.glob(pattern)
        pattern = f"/home/honney/mount/Media/tmp/serien/high_quality/Black Bullet (Season 1)/Season 01/Black Bullet - S01E{i:02}.mkv"
        Quality_Video = glob.glob(pattern)

        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))

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()