from __future__ import annotations

import argparse
import json
import subprocess
import sys
import tempfile
import time
import wave
from pathlib import Path


COMBINATIONS = (
    ("small_vad_gpu_1beam", "small", "GPU", "vad"),
    ("small_context_gpu_1beam", "small", "GPU", "continuous"),
    ("small_context_cpu_5beam", "small", "CPU", "continuous"),
    ("medium_current_baseline", "medium", "GPU", "vad"),
)


def arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Read-only Small/Medium audio A/B diagnostics")
    parser.add_argument("--data-dir", required=True, type=Path)
    parser.add_argument("--media", required=True, type=Path, action="append")
    parser.add_argument("--windows", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--ffmpeg", default="ffmpeg")
    parser.add_argument("--force", action="store_true")
    return parser.parse_args()


def worker_run(data_dir: Path, wav: Path, window: dict, combination: tuple[str, ...]) -> dict:
    name, model, device, segmentation = combination
    with wave.open(str(wav), "rb") as handle:
        sample_rate = handle.getframerate()
    start_ms = int(window["start_ms"])
    duration_ms = int(window["end_ms"]) - start_ms
    command = [
        sys.executable, "-m", "imagefind.audio_worker",
        "--data-dir", str(data_dir), "--wav", str(wav),
        "--device", device, "--model-variant", model,
        "--segmentation", segmentation, "--quality-profile", "accuracy",
        "--language-policy", str(window.get("language_policy", "zh_priority")),
        "--chunk-seconds", str(max(15, min(30, (duration_ms + 999) // 1000))),
        "--overlap-seconds", "0",
        "--start-frame", str(start_ms * sample_rate // 1000),
        "--max-chunks", "1", "--cpu-threads", "2",
    ]
    started = time.monotonic()
    process = subprocess.run(command, text=True, capture_output=True, encoding="utf-8", errors="replace")
    elapsed = time.monotonic() - started
    events = []
    for line in process.stdout.splitlines():
        try:
            events.append(json.loads(line))
        except json.JSONDecodeError:
            pass
    return {
        "combination": name,
        "model": model,
        "device": device,
        "beam": 1 if device == "GPU" else 5,
        "segmentation": segmentation,
        "elapsed_seconds": round(elapsed, 3),
        "window_seconds": round(duration_ms / 1000, 3),
        "rtf": round(elapsed / max(0.001, duration_ms / 1000), 3),
        "exit_code": process.returncode,
        "events": events,
        "stderr_tail": process.stderr[-2000:],
    }


def main() -> int:
    args = arguments()
    if args.output.exists() and not args.force:
        raise SystemExit("output exists; pass --force to replace this diagnostic artifact")
    windows = json.loads(args.windows.read_text(encoding="utf-8"))
    artifact = {
        "schema_version": 1,
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "database_writes": False,
        "runs": [],
    }
    with tempfile.TemporaryDirectory(prefix="imagefind-audio-ab-") as work:
        for media in args.media:
            wav = Path(work) / f"{media.stem}.wav"
            subprocess.run(
                [
                    args.ffmpeg, "-nostdin", "-y", "-i", str(media), "-vn",
                    "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", str(wav),
                ],
                check=True,
                capture_output=True,
            )
            selected = windows.get(media.name)
            if not isinstance(selected, list) or len(selected) != 8:
                raise SystemExit(f"{media.name}: windows JSON must contain exactly 8 labelled windows")
            for window in selected:
                for combination in COMBINATIONS:
                    result = worker_run(args.data_dir, wav, window, combination)
                    result.update({"media": media.name, "window": window})
                    artifact["runs"].append(result)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(artifact, ensure_ascii=False, indent=2), encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())