#!/usr/bin/env python3

import subprocess
import json
import sys
import shutil
from pathlib import Path

def run(cmd):
    return subprocess.run(cmd, capture_output=True, text=True, check=True)


def probe_video(path):
    cmd = [
        "ffprobe",
        "-v", "error",
        "-print_format", "json",
        "-show_streams",
        "-show_format",
        path
    ]

    data = json.loads(run(cmd).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 vaapi_cut(inp, outp, start, end, codec="h264_vaapi"):
    encoder = codec

    cmd = [
        "ffmpeg",
        "-y",
        "-vaapi_device", "/dev/dri/renderD128",

        "-ss", str(start),
        "-to", str(end),

        "-i", inp,
        "-map", "0",

        "-vf", "format=nv12,hwupload",

        "-c:v", encoder,
        "-qp", "20",

        "-c:a", "copy",
        "-c:s", "copy",
        "-c:d", "copy",
        "-c:t", "copy",

        outp
    ]

    subprocess.run(cmd, check=True)

def extract_audio(output_mkv, out_audio):
    # Extract the first (and only) audio track
    run([
        "mkvextract",
        "tracks",
        output_mkv,
        f"1:{out_audio}"  # track IDs are usually 0 for video, 1 for audio
    ])

def mux_audio_into_copy(inp2, audio_file, output_dir):
    output_dir = Path(output_dir)
    output_dir.mkdir(exist_ok=True)

    copied_video = output_dir / Path(inp2).name
    shutil.copy2(inp2, copied_video)

    final_output = output_dir / "final_with_audio.mkv"

    run([
        "mkvmerge",
        "-o", str(final_output),
        str(copied_video),
        "--language", "0:deu",  # assign German language
        str(audio_file)
    ])

    return final_output

def main():
    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 S01E001.mkv"
    inp2 = f"/home/honney/projects/mkv/Anohana - The Flower We Saw That Day - S01 (BD 1080p HEVC Opus) [derpie]/[derpie] Anohana - The Flower We Saw That Day - S01E01 (BD 1080p HEVC Opus) [Dual Audio].mkv"
    
    
    outp = "./output.mkv"
    n1 = 24

    fps, _, _, vstart = probe_video(inp)
    _, duration, _, _ = probe_video(inp2)

    start = vstart + (n1 / fps)
    end = duration + start

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

    print(f"fps: {fps}")
    print(f"cutting from frame {n1} to {end*fps}")
    print(f"time range: {start:.6f}s -> {end:.6f}s")

    print("Trying stream copy...")
    vaapi_cut(inp, outp, start, end)

    output_dir = "./output_dir"
    audio_file = Path(output_dir) / "german_audio.aac"

    extract_audio(outp, audio_file)
    
    final_file = mux_audio_into_copy(inp2, audio_file, output_dir)

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

if __name__ == "__main__":
    main()