#!/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
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):
#     """
#     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 get_video_signal(path):
    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, original_fps, target_fps):
    duration = len(sig) / original_fps
    target_length = int(duration * target_fps)

    return signal.resample(sig, target_length)


def find_offset(video_a, video_b):

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

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

    # 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/Black Bullet/Season 01"
    # for i in range(1, 14):
    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))
        # 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()