from pathlib import Path
import subprocess
import numpy as np
from scipy.signal import correlate
import tempfile
import soundfile as sf
import matplotlib.pyplot as plt
import subprocess
import sys
import os
import glob
import shutil

DEBUG = False

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 extract_audio(file: Path, duration: float):
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)

    subprocess.run([
        "ffmpeg",
        "-y",
        "-i", str(file),
        "-map", "0:a:0",        # first audio stream
        "-ac", "1",             # mono
        "-ar", "44100",         # fixed sample rate
        "-t", str(duration),    # limit duration
        tmp.name
    ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    data, sr = sf.read(tmp.name)
    return data, sr


def get_audio_offset(file1: Path, file2: Path, important_time: float):
    sig1, sr1 = extract_audio(file1, important_time)
    sig2, sr2 = extract_audio(file2, important_time)

    if sr1 != sr2:
        raise ValueError("Sample rates do not match")

    # normalize signals
    sig1 = sig1 - np.mean(sig1)
    sig2 = sig2 - np.mean(sig2)

    corr = correlate(sig1, sig2, mode="full")
    lag = np.argmax(corr) - (len(sig2) - 1)

    offset_seconds = lag / sr1
    return 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 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")
    
    offset = get_audio_offset(Quality_Video, German_Video, 30)
    print(f"Offset: {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 os.path.isfile(quality_audio_file1):
        os.remove(quality_audio_file1)
    if os.path.isfile(quality_audio_file2):
        os.remove(quality_audio_file2)

    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)

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

    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/Maid Sama All/test"
    # for i in range(1, 14):
    for i in [7, 8, 10, 11, 13, 16, 17, 19, 20, 22, 25]:
        # --- German_Videouts ---
        pattern = f"/home/honney/mount/Media/tmp/serien/Maid Sama All/Maid Sama/S01/{i:02}.mkv"
        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)

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