#!/usr/bin/env python3
"""Generate UI SFX via local TangoFluxService (see D:/UD/AI/XModel/TangoFlux_AI调用指南.md)."""
from __future__ import annotations

import argparse
import json
import shutil
import time
from pathlib import Path

import requests

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_MANIFEST = ROOT / "Doc/Design/Sound/UiSfx/manifest.json"


def wait_ready(base: str, timeout: float = 1800) -> dict:
    deadline = time.time() + timeout
    last = {}
    while time.time() < deadline:
        try:
            last = requests.get(f"{base}/api/status", timeout=30).json()
            if last.get("ready"):
                return last
            if not last.get("loaded") and not last.get("busy"):
                try:
                    requests.post(f"{base}/api/warmup", timeout=10)
                except requests.Timeout:
                    pass
        except requests.RequestException as exc:
            last = {"error": str(exc)}
        time.sleep(2)
    raise TimeoutError(f"TangoFlux not ready: {last}")


def submit(base: str, item: dict, defaults: dict) -> dict:
    body = {
        "prompt": item["prompt"],
        "duration": float(item.get("duration", defaults.get("duration", 2))),
        "steps": int(item.get("steps", defaults.get("steps", 50))),
        "guidance_scale": float(item.get("guidance_scale", defaults.get("guidance_scale", 4.5))),
        "seed": int(item.get("seed", 0)),
        "sync": False,
    }
    while True:
        r = requests.post(f"{base}/api/generate", json=body, timeout=60)
        if r.status_code == 429:
            time.sleep(3)
            continue
        r.raise_for_status()
        return r.json()


def wait_job(base: str, job_id: str, timeout: float = 900) -> dict:
    deadline = time.time() + timeout
    while time.time() < deadline:
        j = requests.get(f"{base}/api/jobs/{job_id}", timeout=30).json()
        status = j.get("status")
        if status == "succeeded":
            return j
        if status == "failed":
            raise RuntimeError(j.get("error") or f"job failed: {job_id}")
        time.sleep(0.8)
    raise TimeoutError(job_id)


def download(base: str, job: dict, out_path: Path) -> Path:
    url = job.get("url") or job.get("download_url") or f"/api/jobs/{job['id']}/audio"
    if url.startswith("/"):
        url = base + url
    audio = requests.get(url, timeout=120)
    audio.raise_for_status()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_bytes(audio.content)
    return out_path


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
    ap.add_argument("--base", default="")
    ap.add_argument("--force", action="store_true", help="Regenerate even if WAV exists")
    ap.add_argument("--max-inflight", type=int, default=3)
    args = ap.parse_args()

    manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
    base = (args.base or manifest.get("service") or "http://127.0.0.1:8012").rstrip("/")
    out_dir = ROOT / manifest["outputDir"]
    design_dir = ROOT / manifest["designDir"]
    bundle_dir = ROOT / manifest["bundleDir"]
    defaults = manifest.get("defaults") or {}
    out_dir.mkdir(parents=True, exist_ok=True)
    design_dir.mkdir(parents=True, exist_ok=True)
    bundle_dir.mkdir(parents=True, exist_ok=True)

    status = wait_ready(base)
    print(json.dumps({"ready": True, "status": status}, ensure_ascii=False))

    results = []
    pending: list[tuple[dict, dict]] = []

    def flush_one() -> None:
        item, job = pending.pop(0)
        done = wait_job(base, job["id"])
        wav = out_dir / f"{item['id']}.wav"
        download(base, done, wav)
        dest = bundle_dir / f"{item['id']}.wav"
        shutil.copy2(wav, dest)
        meta = {
            "id": item["id"],
            "jobId": done["id"],
            "prompt": item["prompt"],
            "duration": item.get("duration", defaults.get("duration")),
            "steps": item.get("steps", defaults.get("steps")),
            "guidance_scale": item.get("guidance_scale", defaults.get("guidance_scale")),
            "seed": item.get("seed", 0),
            "bytes": wav.stat().st_size,
            "outputWav": str(wav.relative_to(ROOT)).replace("\\", "/"),
            "bundleWav": str(dest.relative_to(ROOT)).replace("\\", "/"),
            "status": done.get("status"),
        }
        (out_dir / f"{item['id']}.json").write_text(
            json.dumps({**meta, "serverJob": done}, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        results.append(meta)
        print(json.dumps({"done": item["id"], "bytes": meta["bytes"]}, ensure_ascii=False))

    for item in manifest["items"]:
        wav = out_dir / f"{item['id']}.wav"
        if wav.exists() and not args.force:
            dest = bundle_dir / f"{item['id']}.wav"
            if not dest.exists() or args.force:
                shutil.copy2(wav, dest)
            results.append({
                "id": item["id"],
                "skipped": True,
                "outputWav": str(wav.relative_to(ROOT)).replace("\\", "/"),
                "bundleWav": str(dest.relative_to(ROOT)).replace("\\", "/"),
                "bytes": wav.stat().st_size,
            })
            print(json.dumps({"skip": item["id"]}, ensure_ascii=False))
            continue

        while len(pending) >= args.max_inflight:
            flush_one()
        job = submit(base, item, defaults)
        pending.append((item, job))
        print(json.dumps({"queued": item["id"], "jobId": job["id"]}, ensure_ascii=False))

    while pending:
        flush_one()

    report = {
        "batchId": manifest.get("batchId"),
        "base": base,
        "generatedUtc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "count": len(results),
        "items": results,
    }
    (out_dir / "batch_report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    (design_dir / "generation.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(json.dumps({"batchComplete": True, "count": len(results)}, ensure_ascii=False))


if __name__ == "__main__":
    main()