#!/usr/bin/env python3
"""UI SFX v2: generate 1s TangoFlux one-shots per seed, trim to button-length, auto-pick a candidate.

Usage:
    python Tools/generate_ui_sfx_v2.py              # generate missing candidates, process, install picks
    python Tools/generate_ui_sfx_v2.py --process    # reprocess existing candidates only (no service needed)
    python Tools/generate_ui_sfx_v2.py --pick Ui_Click=52003   # override the automatic pick
"""
from __future__ import annotations

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

import numpy as np

from generate_ui_sfx import download, submit, wait_job, wait_ready

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


def read_wav(path: Path) -> tuple[np.ndarray, int]:
    with wave.open(str(path)) as w:
        if w.getsampwidth() != 2:
            raise ValueError(f"{path}: expected 16-bit PCM")
        data = np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16).reshape(-1, w.getnchannels())
        return data.astype(np.float64) / 32768.0, w.getframerate()


def write_wav(path: Path, data: np.ndarray, rate: int) -> None:
    pcm = (np.clip(data, -1, 1) * 32767).round().astype(np.int16)
    path.parent.mkdir(parents=True, exist_ok=True)
    with wave.open(str(path), "wb") as w:
        w.setnchannels(pcm.shape[1])
        w.setsampwidth(2)
        w.setframerate(rate)
        w.writeframes(pcm.tobytes())


def process(data: np.ndarray, rate: int, target: float, peak_db: float, fade_in_ms: float) -> tuple[np.ndarray, dict]:
    """Trim leading silence, cap at target with an exponential-ish fade, normalize peak. Returns clip and score metrics."""
    level = np.abs(data).max(axis=1)
    peak = level.max()
    if peak < 0.1:
        # Near-silent generations would only amplify the noise floor once normalized.
        return data[:1], {"onsetSec": 0.0, "kept": 0.0, "score": -1.0, "rawPeak": round(float(peak), 4), "silent": True}
    onset = int(np.argmax(level > peak * 0.08))
    onset = max(0, onset - int(rate * 0.003))
    energy = level ** 2
    n = int(target * rate)
    kept = float(energy[onset:onset + n].sum() / energy.sum())
    clip = data[onset:onset + n].copy()
    fade_in = min(len(clip), max(1, int(rate * fade_in_ms / 1000)))
    clip[:fade_in] *= np.linspace(0, 1, fade_in)[:, None]
    # Fade the last 40% to zero so the cap never clicks; squared curve keeps the attack intact.
    tail = max(1, int(len(clip) * 0.4))
    clip[-tail:] *= (np.linspace(1, 0, tail) ** 2)[:, None]
    clip *= (10 ** (peak_db / 20)) / max(1e-6, np.abs(clip).max())
    onset_sec = onset / rate
    # Prefer energy already inside the window (less truncation) and an immediate attack.
    score = kept - 2.0 * max(0.0, onset_sec - 0.02)
    return clip, {"onsetSec": round(onset_sec, 4), "kept": round(kept, 4), "score": round(score, 4), "rawPeak": round(float(peak), 4), "silent": False}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
    ap.add_argument("--process", action="store_true", help="Skip generation; reprocess existing candidates")
    ap.add_argument("--pick", action="append", default=[], help="Override pick, e.g. Ui_Click=52003")
    ap.add_argument("--max-inflight", type=int, default=3)
    args = ap.parse_args()

    manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
    defaults = manifest["defaults"]
    out_dir = ROOT / manifest["outputDir"]
    bundle_dir = ROOT / manifest["bundleDir"]
    design_dir = ROOT / manifest["designDir"]
    picks = dict(p.split("=", 1) for p in args.pick)
    base = manifest["service"].rstrip("/")

    if not args.process:
        wait_ready(base)
        pending: list[tuple[Path, dict, dict]] = []

        def flush_one() -> None:
            path, body, job = pending.pop(0)
            done = wait_job(base, job["id"])
            download(base, done, path)
            path.with_suffix(".json").write_text(json.dumps({**body, "serverJob": done}, ensure_ascii=False, indent=2), encoding="utf-8")
            print(json.dumps({"done": path.name}))

        for item in manifest["items"]:
            for seed in item.get("seeds", defaults["seeds"]):
                path = out_dir / "raw" / f"{item['id']}_{seed}.wav"
                if path.exists():
                    continue
                while len(pending) >= args.max_inflight:
                    flush_one()
                body = {"prompt": item["prompt"] + defaults["suffix"], "duration": defaults["duration"], "seed": seed}
                job = submit(base, body, defaults)
                pending.append((path, body, job))
        while pending:
            flush_one()

    report = []
    for item in manifest["items"]:
        candidates = []
        for seed in item.get("seeds", defaults["seeds"]):
            raw = out_dir / "raw" / f"{item['id']}_{seed}.wav"
            data, rate = read_wav(raw)
            clip, metrics = process(data, rate, item["targetSec"], item["peakDb"], defaults["fadeInMs"])
            trimmed = out_dir / "trimmed" / f"{item['id']}_{seed}.wav"
            write_wav(trimmed, clip, rate)
            candidates.append({"seed": seed, "raw": str(raw.relative_to(ROOT)).replace("\\", "/"),
                               "trimmed": str(trimmed.relative_to(ROOT)).replace("\\", "/"),
                               "lengthSec": round(len(clip) / rate, 4), **metrics})
        auto = max(candidates, key=lambda c: c["score"])
        chosen = next((c for c in candidates if str(c["seed"]) == picks.get(item["id"])), auto)
        shutil.copy2(ROOT / chosen["trimmed"], bundle_dir / f"{item['id']}.wav")
        report.append({"id": item["id"], "targetSec": item["targetSec"], "peakDb": item["peakDb"],
                       "prompt": item["prompt"] + defaults["suffix"], "pickedSeed": chosen["seed"],
                       "pickedBy": "manual" if chosen is not auto else "auto", "candidates": candidates})
        print(json.dumps({"id": item["id"], "seed": chosen["seed"], "kept": chosen["kept"], "onset": chosen["onsetSec"]}))

    result = {"batchId": manifest["batchId"], "generatedUtc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "items": report}
    (out_dir / "batch_report.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    (design_dir / "generation_v2.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")


if __name__ == "__main__":
    main()