"""Multi-reference a2a generation for the AKUSPACE LoRA — one model load, many sources.

ltx_a2a_generate.py pins a single --reference to every prompt, so an N-source
sweep costs N model loads per checkpoint. But validation samples carry their own
conditions[].audio, so one load can cover every (source, prompt) pair. Loading
dominates runtime, so this is the difference between 15 loads and 3.

Input is a TSV manifest (label<TAB>reference_wav<TAB>prompt), which avoids the
shell-quoting traps that bit us with long caption strings.

Outputs land as samples/step_000001_<n>.wav in prompt order; this renames them
to <label>.wav so a 60-clip sweep stays readable.

Trigger note: matches ltx_a2a_generate.py — AKUSPACE is prepended here because
validation_runner does not add it, while training captions had it.
"""

import argparse
import copy
import subprocess
import sys
from pathlib import Path

import yaml

DEFAULT_TRIGGER = "AKUSPACE"
NEVER_SAVE = 10**9


def read_manifest(path: Path) -> list[tuple[str, str, str]]:
    rows = []
    for n, line in enumerate(path.read_text().splitlines(), 1):
        line = line.rstrip("\n")
        if not line.strip() or line.lstrip().startswith("#"):
            continue
        parts = line.split("\t")
        if len(parts) != 3:
            raise SystemExit(f"{path}:{n}: expected 3 tab-separated fields, got {len(parts)}")
        label, ref, prompt = (p.strip() for p in parts)
        if not Path(ref).exists():
            raise SystemExit(f"{path}:{n}: reference not found: {ref}")
        rows.append((label, ref, prompt))
    if not rows:
        raise SystemExit(f"{path}: no rows")
    return rows


def build_config(args, rows) -> dict:
    cfg = yaml.safe_load(Path(args.base_config).read_text())
    cfg["model"]["load_checkpoint"] = str(Path(args.checkpoint).resolve())
    cfg.setdefault("checkpoints", {})["no_resume"] = True
    cfg["checkpoints"]["interval"] = NEVER_SAVE
    cfg["optimization"]["steps"] = 1

    samples = []
    for _label, ref, prompt in rows:
        text = prompt if args.no_trigger else f"{args.trigger} {prompt}"
        samples.append(
            {
                "prompt": text,
                "conditions": [{"type": "reference", "audio": str(Path(ref).resolve())}],
            }
        )

    validation = copy.deepcopy(cfg.get("validation", {}))
    validation["samples"] = samples
    validation["interval"] = 1
    if args.inference_steps:
        validation["inference_steps"] = args.inference_steps
    if args.seed is not None:
        validation["seed"] = args.seed
    cfg["validation"] = validation
    cfg["output_dir"] = str(Path(args.output_dir).resolve())
    return cfg


def rename_outputs(out: Path, rows) -> int:
    """samples/step_000001_<n>.wav -> <label>.wav, n is 1-based prompt order."""
    sample_dir = out / "samples"
    renamed = 0
    for idx, (label, _ref, _prompt) in enumerate(rows, 1):
        matches = sorted(sample_dir.glob(f"step_*_{idx}.wav"))
        if not matches:
            print(f"  MISSING output for #{idx} {label}")
            continue
        dest = out / f"{label}.wav"
        dest.write_bytes(matches[-1].read_bytes())
        renamed += 1
    return renamed


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--manifest", required=True, help="TSV: label<TAB>reference<TAB>prompt")
    p.add_argument("--output-dir", required=True)
    p.add_argument("--base-config", required=True)
    p.add_argument("--trigger", default=DEFAULT_TRIGGER)
    p.add_argument("--no-trigger", action="store_true")
    p.add_argument("--inference-steps", type=int, default=None)
    p.add_argument("--seed", type=int, default=None)
    p.add_argument("--ltx-repo", default="/workspace/LTX-2.5-repo")
    p.add_argument("--dry-run", action="store_true")
    args = p.parse_args()

    for path in (args.checkpoint, args.base_config, args.manifest):
        if not Path(path).exists():
            print(f"missing: {path}")
            return 2

    rows = read_manifest(Path(args.manifest))
    cfg = build_config(args, rows)
    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)
    cfg_path = out / "generate_config.yaml"
    cfg_path.write_text(yaml.safe_dump(cfg, sort_keys=False))

    refs = sorted({r[1] for r in rows})
    print(f"config     : {cfg_path}")
    print(f"checkpoint : {cfg['model']['load_checkpoint']}")
    print(f"prompts    : {len(rows)} across {len(refs)} reference(s), ONE model load")
    print(f"trigger    : {'off' if args.no_trigger else args.trigger}")
    for i, (label, ref, _prompt) in enumerate(rows, 1):
        print(f"   {i:3d}. {label:38s} <- {Path(ref).name}")

    if args.dry_run:
        return 0

    cmd = ["uv", "run", "python", "scripts/train.py", str(cfg_path)]
    cwd = Path(args.ltx_repo) / "packages" / "ltx-trainer"
    print(f"\nrunning: {' '.join(cmd)}  (cwd={cwd})")
    result = subprocess.run(cmd, cwd=cwd)
    if result.returncode == 0:
        n = rename_outputs(out, rows)
        print(f"\nlabelled {n}/{len(rows)} clips in {out}")
    return result.returncode


if __name__ == "__main__":
    sys.exit(main())