"""Export one scenario's (.npz + per-turn audio) into browser-playable form.

Produces, in data/viewer/:
    {scenario_id}.json  — {fps, num_frames, turns, blendshapes[T][52], names[52]}
    {scenario_id}.mp3   — concatenated turn audio, aligned to blendshape timeline

Usage:
    python -m scripts.compiler.export_for_viewer <scenario_id>
"""
from __future__ import annotations

import argparse
import json
import subprocess
from pathlib import Path

import numpy as np

from .constants import ARKIT_52_NAMES

PROJECT_ROOT = Path(__file__).resolve().parents[2]
NPZ_DIR = PROJECT_ROOT / "data" / "v3_training"
AUDIO_DIR = PROJECT_ROOT / "data" / "audio_preview"
SCENARIO_FILES = [
    PROJECT_ROOT / "data" / "emotion" / "seed_train_final.jsonl",
    PROJECT_ROOT / "data" / "emotion" / "seed_val.jsonl",
    PROJECT_ROOT / "data" / "emotion" / "seed_test.jsonl",
]
OUT_DIR = PROJECT_ROOT / "data" / "viewer"
FPS = 30


def find_scenario(scenario_id: str):
    for p in SCENARIO_FILES:
        if not p.exists():
            continue
        with p.open() as f:
            for line in f:
                scen = json.loads(line)
                if scen["scenario_id"] == scenario_id:
                    return scen
    return None


def find_audio(scenario_id: str, turn_idx: int, emotion: str) -> Path | None:
    exact = AUDIO_DIR / f"{scenario_id}_t{turn_idx}_{emotion}.mp3"
    if exact.exists() and exact.stat().st_size > 1000:
        return exact
    matches = [m for m in AUDIO_DIR.glob(f"{scenario_id}_t{turn_idx}_*.mp3")
               if m.stat().st_size > 1000]
    return matches[0] if matches else None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("scenario_id")
    args = ap.parse_args()

    OUT_DIR.mkdir(parents=True, exist_ok=True)

    npz_path = NPZ_DIR / f"{args.scenario_id}.npz"
    if not npz_path.exists():
        raise SystemExit(f"Not found: {npz_path}")

    scen = find_scenario(args.scenario_id)
    if not scen:
        raise SystemExit(f"Scenario not found in seeds: {args.scenario_id}")

    data = np.load(npz_path)
    target = data["target"]  # (T, 52)
    print(f"[load] blendshapes shape: {target.shape}")

    # Collect per-turn audio in scenario order
    audio_paths = []
    turns_meta = []
    for ti, turn in enumerate(scen["turns"]):
        if not turn.get("text", "").strip():
            continue
        emo = turn.get("emotion", "neutral")
        p = find_audio(args.scenario_id, ti, emo)
        if p is None:
            raise SystemExit(f"Missing audio for turn {ti} ({emo})")
        audio_paths.append(p)
        turns_meta.append({
            "turn_idx": ti,
            "emotion": emo,
            "vad": turn.get("vad"),
            "text": turn.get("text"),
            "speaker": turn.get("speaker"),
        })
    print(f"[audio] {len(audio_paths)} turns to concatenate")

    # Concat audio via ffmpeg concat demuxer
    audio_out = OUT_DIR / f"{args.scenario_id}.mp3"
    list_file = OUT_DIR / f"{args.scenario_id}.concat.txt"
    with list_file.open("w") as f:
        for p in audio_paths:
            f.write(f"file '{p.resolve()}'\n")
    subprocess.run(
        [
            "ffmpeg", "-y", "-loglevel", "error",
            "-f", "concat", "-safe", "0",
            "-i", str(list_file),
            "-c:a", "libmp3lame", "-b:a", "128k",
            str(audio_out),
        ],
        check=True,
    )
    list_file.unlink()

    # Export JSON (round to 4 decimals to keep file small)
    json_out = OUT_DIR / f"{args.scenario_id}.json"
    payload = {
        "scenario_id": args.scenario_id,
        "fps": FPS,
        "num_frames": int(target.shape[0]),
        "names": ARKIT_52_NAMES,
        "turns": turns_meta,
        "blendshapes": np.round(target, 4).tolist(),
    }
    with json_out.open("w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False)

    size_kb = json_out.stat().st_size // 1024
    print(f"[done]")
    print(f"  JSON:  {json_out}  ({size_kb} KB)")
    print(f"  Audio: {audio_out}")
    print(f"\nOpen in browser:")
    print(f"  tools/blendshape-player.html?scenario={args.scenario_id}")


if __name__ == "__main__":
    main()