#!/usr/bin/env python3
"""Apply sync: mux file A video + file A audio + delayed file B audio into output.

Usage:
    python apply_sync.py --config config.toml --lag 0.042
    python apply_sync.py --a video_a.mkv --b video_b.mkv --lag -0.015 --out output/synced.mkv
"""
import argparse
import json
import os
import shutil
import subprocess
import sys
import tomllib


def probe_audio_streams(path: str) -> list[dict]:
    """Probe all audio streams in a file."""
    if shutil.which("ffprobe") is None:
        return []
    try:
        r = subprocess.run(
            ["ffprobe", "-v", "error", "-select_streams", "a",
             "-show_entries",
             "stream=index,codec_name,channels,sample_rate,bit_rate,bits_per_sample,"
             "duration,start_time,start_pts,duration_ts,"
             "stream_disposition=forced,default:"
             "stream_tags=language,title",
             "-of", "json", path],
            capture_output=True, text=True, timeout=10
        )
        if r.returncode != 0:
            return []
        data = json.loads(r.stdout)
        streams = []
        for s in data.get("streams", []):
            tags = {k.lower(): v for k, v in (s.get("tags") or {}).items()}
            dispo = s.get("disposition", {})
            sr = s.get("sample_rate")
            lang_raw = tags.get("language", "und")
            streams.append({
                "index": s.get("index", 0),
                "codec": s.get("codec_name", ""),
                "channels": s.get("channels", 0),
                "sample_rate": int(sr) if sr else 0,
                "language": lang_raw if lang_raw else "und",
                "title": tags.get("title", ""),
                "is_default": dispo.get("default", 0) == 1,
                "is_forced": dispo.get("forced", 0) == 1,
                "start_time": float(s.get("start_time") or 0),
                "start_pts": int(s.get("start_pts") or 0),
                "duration_ts": int(s.get("duration_ts") or 0),
                "duration": float(s.get("duration") or 0),
            })
        return streams
    except Exception:
        return []


def _codec_args(codec: str, out_idx: int) -> list[str]:
    """Return ffmpeg codec args for a given audio codec string."""
    lossless = {"flac", "alac", "pcm_s16le", "pcm_s24le", "pcm_s32le", "pcm_f32le", "pcm_f64le"}
    if codec in lossless:
        return [f"-c:a:{out_idx}", codec]
    # lossy codecs: use original where possible
    codec_map = {
        "aac": "aac", "mp3": "libmp3lame", "opus": "libopus",
        "vorbis": "libvorbis", "ac3": "ac3", "eac3": "eac3",
    }
    enc = codec_map.get(codec, "aac")
    args = [f"-c:a:{out_idx}", enc]
    if enc != codec or codec not in lossless:
        args += [f"-b:a:{out_idx}", "192k"]
    return args


def build_ffmpeg_cmd(file_a: str, file_b: str, lag_sec: float, out_path: str) -> list[str]:
    """Build ffmpeg command to mux A video + A audio + delayed B audio streams."""
    b_streams = probe_audio_streams(file_b)

    lag_ms = lag_sec * 1000

    # Check if B has a single stream with unspecified language
    set_german = (len(b_streams) == 1 and b_streams[0]["language"] in ("und", "unknown", ""))

    # Build filter_complex for B audio streams
    # Positive lag (B behind A): adelay adds silence at start
    # Negative lag (B ahead of A): atrim cuts the start, asetpts resets timestamps
    filters = []
    for i, s in enumerate(b_streams):
        if lag_ms >= 0:
            filters.append(f"[1:a:{i}]adelay={lag_ms:.0f}|{lag_ms:.0f}:all=1[d{i}]")
        else:
            trim_sec = -lag_sec
            filters.append(f"[1:a:{i}]atrim=start={trim_sec:.6f},asetpts=N/SR/TB[d{i}]")

    cmd = ["ffmpeg", "-y", "-i", file_a, "-i", file_b]

    if filters:
        cmd += ["-filter_complex", ";".join(filters)]

    # Map: video from A, all audio from A, all delayed B audio
    cmd += ["-map", "0:v", "-map", "0:a"]
    for i in range(len(b_streams)):
        cmd += ["-map", f"[d{i}]"]

    # Codecs: copy video + A's audio, re-encode B's filtered audio with same codec
    n_a_audio = len(probe_audio_streams(file_a))
    cmd += ["-c:v", "copy"]
    for i in range(n_a_audio):
        cmd += [f"-c:a:{i}", "copy"]
    for i, s in enumerate(b_streams):
        cmd += _codec_args(s["codec"], n_a_audio + i)

    # Language metadata for B audio streams (they become the last N audio streams in output)
    for i, s in enumerate(b_streams):
        out_audio_idx = n_a_audio + i  # 0-based index in output audio streams
        if set_german and i == 0:
            cmd += [f"-metadata:s:a:{out_audio_idx}", "language=ger"]
        elif s["language"] not in ("und", "unknown", ""):
            cmd += [f"-metadata:s:a:{out_audio_idx}", f"language={s['language']}"]
        if s["title"]:
            cmd += [f"-metadata:s:a:{out_audio_idx}", f"title={s['title']}"]
        if s["is_forced"]:
            cmd += [f"-disposition:a:{out_audio_idx}", "forced"]

    cmd.append(out_path)
    return cmd


def main():
    ap = argparse.ArgumentParser(description="Apply sync delay to mux audio streams")
    ap.add_argument("--config", default=None, help="config.toml path")
    ap.add_argument("--a", default=None, help="Video file A (base)")
    ap.add_argument("--b", default=None, help="Video file B (delayed audio source)")
    ap.add_argument("--lag", type=float, required=True, help="Lag in seconds (B relative to A, positive = B delayed)")
    ap.add_argument("--out", default=None, help="Output file path (overrides config)")
    args = ap.parse_args()

    file_a = args.a
    file_b = args.b
    out_path = args.out

    if args.config:
        with open(args.config, "rb") as f:
            cfg = tomllib.load(f)
        ext = cfg.get("extract", {})
        sync = cfg.get("sync", {})
        file_a = file_a or ext.get("file_a")
        file_b = file_b or ext.get("file_b")
        if not out_path:
            out_dir = sync.get("out_dir", "output")
            stem = os.path.splitext(os.path.basename(file_a))[0]
            out_path = os.path.join(out_dir, f"{stem}_synced.mkv")

    if not file_a or not file_b:
        print("ERROR: need --a/--b or --config with file_a/file_b", file=sys.stderr)
        sys.exit(1)
    if not os.path.exists(file_a):
        print(f"ERROR: file A not found: {file_a}", file=sys.stderr); sys.exit(1)
    if not os.path.exists(file_b):
        print(f"ERROR: file B not found: {file_b}", file=sys.stderr); sys.exit(1)

    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)

    b_streams = probe_audio_streams(file_b)
    print(f"File A: {file_a}")
    print(f"File B: {file_b}")
    print(f"B audio streams: {len(b_streams)}")
    for i, s in enumerate(b_streams):
        lang = s['language']
        title = f' "{s["title"]}"' if s["title"] else ""
        print(f"  [{i}] {lang}{title} {s['codec'].upper()} {s['channels']}ch")
    print(f"Lag: {args.lag:+.6f}s ({args.lag*1000:+.0f}ms)")
    print(f"Output: {out_path}")

    cmd = build_ffmpeg_cmd(file_a, file_b, args.lag, out_path)
    print(f"\nRunning: {' '.join(cmd)}")
    r = subprocess.run(cmd)
    if r.returncode != 0:
        print(f"ERROR: ffmpeg failed with code {r.returncode}", file=sys.stderr)
        sys.exit(r.returncode)

    print(f"\nSynced output: {out_path}")

    # Auto-launch verify_sync in full mode
    verify_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "verify_sync.py")
    if os.path.exists(verify_script):
        verify_cmd = [sys.executable, verify_script, "--config", args.config or "config.toml",
                      "--lag", f"{args.lag:.6f}", "--mode", "full"]
        print(f"Launching verify: {' '.join(verify_cmd)}")
        subprocess.Popen(verify_cmd)

    sys.exit(0)


if __name__ == "__main__":
    main()