#!/usr/bin/env python3
"""
Assemble the final MP4s: silent video track + MiniMax narration + subtitles.

Each narration segment was recorded as its own audio file, generated by
MiniMax with a female voice for the language of the track.  The segments are
placed at the moment their beat starts on screen, so the voice, the caption
and the terminal action always agree.

    ./scripts/screencast/assemble.py --phase 04-roteamento --lang en
"""
import argparse
import json
import os
import shutil
import subprocess
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
REPO = os.path.dirname(os.path.dirname(HERE))
CAPS = os.path.join(REPO, "out", "screencast", "caps")
VO = os.path.join(REPO, "out", "screencast", "vo")
OUT = os.path.join(REPO, "out", "screencast")
ART = os.path.join(OUT, "art")

# which phases draw the lab diagram over the terminal
DIAGRAM_PHASES = {"00-ambiente", "04-roteamento", "05-balanceamento"}
# phases that show the module/host instead of the firewall VM
HOST_PHASES = {"01-implementacao"}

PHASE_ORDER = ["00-ambiente", "01-implementacao", "02-ativacao", "03-regras",
               "04-roteamento", "05-balanceamento", "06-qos"]


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


def probe(path):
    out = subprocess.run(["ffprobe", "-v", "error", "-show_entries",
                          "format=duration", "-of", "csv=p=0", path],
                         capture_output=True, text=True).stdout.strip()
    return float(out or 0)


def audio_duration(path):
    return probe(path)


def make_video(phase, lang):
    """Render the silent video (terminal + diagram + burned-in subtitles)."""
    cast = os.path.join(CAPS, f"{phase}.cast")
    nar = os.path.join(HERE, "narration", f"{phase}.{lang}.json")
    if not os.path.exists(cast):
        raise SystemExit(f"[assemble] missing capture: {cast}")
    art = None
    if os.path.isdir(ART):
        imgs = sorted(f for f in os.listdir(ART) if f.endswith((".jpg", ".png")))
        if imgs:
            art = os.path.join(ART, imgs[0])
    outdir = os.path.join(OUT, lang)
    os.makedirs(outdir, exist_ok=True)
    silent = os.path.join(outdir, f"{phase}.silent.mp4")
    index = PHASE_ORDER.index(phase) + 1
    cmd = [sys.executable, os.path.join(HERE, "render_video.py"),
           "--cast", cast, "--narration", nar, "--out", silent,
           "--lang", lang, "--phase-no", str(index),
           "--phase-total", str(len(PHASE_ORDER))]
    if phase in DIAGRAM_PHASES:
        cmd.append("--diagram")
    if art:
        cmd += ["--art", art]
    run(cmd)
    return silent


def mux_audio(phase, lang, silent):
    """Lay every narration segment on the timeline and add it to the video."""
    meta = json.load(open(os.path.join(CAPS, f"{phase}.json")))
    beats = {b["id"]: b for b in meta["beats"]}
    nar = json.load(open(os.path.join(HERE, "narration", f"{phase}.{lang}.json")))
    vodir = os.path.join(VO, f"{phase}.{lang}")

    segments = []
    for i, seg in enumerate(nar["segments"]):
        beat = beats.get(seg["beat"])
        audio = os.path.join(vodir, f"{i:02d}-{seg['beat']}.mp3")
        if beat is None or not os.path.exists(audio):
            sys.stderr.write(f"[assemble] missing audio: {audio}\n")
            continue
        segments.append((beat["start"], audio))

    tmp = os.path.join(OUT, "_tmp")
    os.makedirs(tmp, exist_ok=True)
    if not segments:
        out = os.path.join(tmp, f"{phase}.{lang}.m4v")
        run(["ffmpeg", "-y", "-loglevel", "error", "-i", silent,
             "-c:v", "copy", "-an", out])
        return out

    vdur = probe(silent)
    inputs = []
    filters = []
    # input 0 is the video; the narration files are inputs 1..N
    for k, (start, audio) in enumerate(segments):
        inputs += ["-i", audio]
        delay = int(start * 1000)
        # pad so every segment reaches the end of the video, then the mix is
        # trimmed to the video length
        filters.append(f"[{k + 1}:a]adelay={delay}|{delay},apad[aud{k}]")
    mix_in = "".join(f"[aud{k}]" for k in range(len(segments)))
    filters.append(f"{mix_in}amix=inputs={len(segments)}:normalize=0,"
                   f"atrim=0:{vdur:.3f},asetpts=N/SR/TB[mixed]")
    cmd = ["ffmpeg", "-y", "-loglevel", "error", "-i", silent] + inputs + \
          ["-filter_complex", ";".join(filters),
           "-map", "0:v", "-map", "[mixed]", "-c:v", "copy", "-c:a", "aac",
           "-b:a", "160k", "-ar", "44100", "-ac", "2",
           os.path.join(tmp, f"{phase}.{lang}.m4v")]
    run(cmd)
    return os.path.join(tmp, f"{phase}.{lang}.m4v")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--phase", required=True)
    ap.add_argument("--lang", required=True, choices=["en", "pt"])
    ap.add_argument("--keep-silent", action="store_true")
    args = ap.parse_args()

    phase, lang = args.phase, args.lang
    silent = make_video(phase, lang)
    muxed = mux_audio(phase, lang, silent)

    outdir = os.path.join(OUT, lang)
    os.makedirs(outdir, exist_ok=True)
    final = os.path.join(outdir, f"{phase}.mp4")
    run(["ffmpeg", "-y", "-loglevel", "error", "-i", muxed,
         "-c", "copy", "-movflags", "+faststart", final])
    if not args.keep_silent:
        os.remove(silent)
    os.remove(muxed)

    dur = probe(final)
    size = os.path.getsize(final) / 1e6
    print(f"[assemble] {final}  {dur:.1f}s  {size:.1f}MB")


if __name__ == "__main__":
    main()