#!/usr/bin/env python3
"""Generate Russian music with the measured YuE2 packed inference profiles."""
import argparse
import dataclasses
import json
import math
import hashlib
import os
import sys
import time

PROCESS_T0 = time.time()  # module import; the interpreter launch itself is not included
from pathlib import Path

SRC = Path(__file__).absolute().parent
ROOT = SRC.parent
sys.path.insert(0, str(SRC / "runtime/kernels/base"))
sys.path.insert(0, str(SRC / "runtime"))
PROFILES = ("reference", "fused", "fast", "memory")


def load_pipeline(*, vae=None, low_memory=False, offload_ar=False, progress=True,
                  profile=None, ode_steps=None, vae_precision=None, fuse_rms_quant=False,
                  capture_sampling=None, kv_cache_dtype="bf16", int8_attention=None):
    import torch
    profile = profile or ("reference" if low_memory else "fast")
    if profile not in PROFILES:
        raise ValueError(f"Unknown profile {profile!r}")
    if low_memory and profile != "reference":
        raise ValueError("Unfused --low-memory requires profile=reference; use profile=memory for staged experts")
    if fuse_rms_quant and profile not in {"fast", "memory"}:
        raise ValueError("RMS/activation fusion requires profile=fast or profile=memory")
    if kv_cache_dtype not in {"bf16", "int8"}:
        raise ValueError("kv_cache_dtype must be bf16 or int8")
    if kv_cache_dtype == "int8" and (profile not in {"fast", "memory"} or not fuse_rms_quant):
        raise ValueError("INT8 KV requires profile=fast or memory and fuse_rms_quant=True")
    if int8_attention is not None and int8_attention not in {"reference", "ordered", "subwarp"}:
        raise ValueError("int8_attention must be reference, ordered, subwarp, or None")
    if int8_attention is not None and kv_cache_dtype != "int8":
        raise ValueError("int8_attention requires kv_cache_dtype=int8")
    int8_attention = (int8_attention or "ordered") if kv_cache_dtype == "int8" else None
    if capture_sampling is not None and type(capture_sampling) is not bool:
        raise ValueError("capture_sampling must be a boolean or None")
    if capture_sampling and profile not in {"fast", "memory"}:
        raise ValueError("Captured sampling requires profile=fast or profile=memory")
    if not torch.cuda.is_available():
        raise RuntimeError("This release requires an NVIDIA CUDA GPU.")
    if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}:
        raise RuntimeError("This runtime ships kernels for SM89 (Ada/RTX 40) and SM120 (Blackwell consumer); other GPUs need a rebuild.")
    if ode_steps is not None and (isinstance(ode_steps, bool) or not isinstance(ode_steps, int) or ode_steps < 1):
        raise ValueError("ode_steps must be a positive integer")
    if vae_precision is not None and vae_precision not in {"fp32", "fp16", "bf16"}:
        raise ValueError("vae_precision must be fp32, fp16, or bf16")
    if vae is None:
        from huggingface_hub import snapshot_download
        lock = json.loads((ROOT / "source-lock.json").read_text())
        vae = snapshot_download("m-a-p/YuE2-Vae", revision=lock["m-a-p/YuE2-Vae"],
                                allow_patterns=["*.json", "*.safetensors", "*.py", "LICENSE", "licenses/*", "THIRD_PARTY_NOTICES.md"])
    from yue2_orbit import OrbitPipeline
    pipe = OrbitPipeline.from_pretrained(str(ROOT), vae=vae, device="cuda",
                                         memory_budget_gib=18, vae_core_frames=512,
                                         offload_ar=offload_ar or profile == "memory", progress=progress)
    pipe.fused = not low_memory
    pipe.trim_cache = True
    pipe.profile = profile
    pipe.fuse_rms_quant = bool(fuse_rms_quant)
    pipe.kv_cache_dtype = kv_cache_dtype
    pipe.int8_attention = int8_attention
    fast = profile in {"fast", "memory"}
    pipe.capture_sampling = fast if capture_sampling is None else capture_sampling
    pipe.generation_config = dataclasses.replace(pipe.generation_config,
                                                 ode_steps=ode_steps if ode_steps is not None else 8 if fast else 32)
    precision = vae_precision or ("fp16" if fast else "fp32")
    pipe.vae_compute_dtype = {"fp32": None, "fp16": torch.float16, "bf16": torch.bfloat16}[precision]
    # Include packed adapters and binaries, not only the upstream yue2 package.
    files = sorted(p for p in SRC.rglob("*") if p.is_file() and
                   "__pycache__" not in p.parts and "kernel-source" not in p.parts and
                   (p.suffix in {".py", ".so", ".whl"} or p.name == "requirements.txt"))
    hashes = {str(p.relative_to(SRC)): hashlib.sha256(p.read_bytes()).hexdigest() for p in files}
    pipe.runtime_bundle_sha256 = hashlib.sha256(json.dumps(hashes, sort_keys=True).encode()).hexdigest()
    return pipe


def _guidance_scale(value):
    value = float(value)
    if not math.isfinite(value) or not 0 <= value <= 20:
        raise argparse.ArgumentTypeError("cfg scale must be finite and in [0,20]")
    return value


class _NvmlPeak:
    """Whole-process NVML peak of this GPU's used memory (20 ms sampling), for the report."""

    def __init__(self):
        import threading
        self.peak = 0; self._stop = threading.Event(); self._thread = None
        try:
            import pynvml
            pynvml.nvmlInit(); self._h = pynvml.nvmlDeviceGetHandleByIndex(0); self._nv = pynvml
        except Exception:
            self._nv = None
            return
        self._thread = threading.Thread(target=self._run, daemon=True); self._thread.start()

    def _run(self):
        while not self._stop.wait(0.02):
            try:
                self.peak = max(self.peak, self._nv.nvmlDeviceGetMemoryInfo(self._h).used)
            except Exception:
                return

    def stop(self):
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=1)
        return self.peak / 2**30 if self._nv is not None else None


def main():
    nvml_peak = _NvmlPeak()
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prompt", default=str(ROOT / "examples/prompts/default-ru.json"))
    parser.add_argument("--output", default="song")
    parser.add_argument("--vae")
    parser.add_argument("--profile", choices=PROFILES)
    parser.add_argument("--ode-steps", type=int)
    parser.add_argument("--cfg-scale", type=_guidance_scale,
                        help="Semantic guidance scale; 1 uses one decode branch")
    parser.add_argument("--vae-precision", choices=("fp32", "fp16", "bf16"))
    parser.add_argument("--allocator", choices=("native", "expandable"))
    parser.add_argument("--low-memory", action="store_true")
    parser.add_argument("--offload-ar", action="store_true")
    parser.add_argument("--fuse-rms-quant", action="store_true",
                        help="Fuse RMSNorm into activation quantization; faster experimental arithmetic")
    parser.add_argument("--kv-cache-dtype", choices=["bf16", "int8"], default="bf16",
                        help="AR KV cache precision; int8 requires --fuse-rms-quant and fast/memory")
    parser.add_argument("--int8-attention", choices=["reference", "ordered", "subwarp"],
                        help="INT8 decode: ordered preserves previous arithmetic (default); subwarp changes rounding; reference uses the previous kernel")
    parser.add_argument("--no-captured-sampling", action="store_true",
                        help="Use the earlier eager sampler for comparison")
    parser.add_argument("--warm-compile", action="store_true",
                        help="keep the torch.compile paths (VAE decoder, NAR pointwise chains); pays off only when one process serves several songs")
    parser.add_argument("--mode", choices=["off", "exact", "fast", "lowmem", "turbo"], default="exact",
                        help="off: reference decode path; exact: same sampled tokens, faster pipeline; fast: + INT8 output heads and NAR KV handoff; lowmem: + INT8 embeddings; turbo: W4A4 output heads")
    args = parser.parse_args()
    allocator = args.allocator or ("expandable" if (args.profile or "fast") in {"fast", "memory"} and not args.low_memory else None)
    if allocator is not None:
        os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" if allocator == "expandable" else "backend:native"
    request = json.loads(Path(args.prompt).read_text())
    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=True)
    with load_pipeline(vae=args.vae, low_memory=args.low_memory, offload_ar=args.offload_ar,
                       profile=args.profile, ode_steps=args.ode_steps, vae_precision=args.vae_precision,
                       fuse_rms_quant=args.fuse_rms_quant, kv_cache_dtype=args.kv_cache_dtype,
                       int8_attention=args.int8_attention,
                       capture_sampling=False if args.no_captured_sampling else None) as pipe:
        mode_flags = []
        if args.mode != "off":
            from fast_ops import apply_mode
            model = pipe._load_model()
            mode_flags = apply_mode(pipe, model, args.mode, warm=args.warm_compile)
        gen_start = time.perf_counter()
        generation = {k: request[k] for k in ("style", "lyrics", "seed", "cot") if k in request}
        if args.cfg_scale is not None:
            generation["cfg_scale"] = args.cfg_scale
        song = pipe(**generation)
        t_save = time.perf_counter()
        song.save_artifacts(out)
        t_artifacts = time.perf_counter() - t_save
        song.save(out / "audio.wav")
        t_wav = time.perf_counter() - t_save - t_artifacts
        stage_peaks = [v.get("peak_gib", 0) for v in song.timing.get("stages", {}).get("memory", {}).values()]
        report = dict(seconds_including_load=time.time() - PROCESS_T0, generation_seconds=time.perf_counter() - gen_start,
                      mode=args.mode, mode_flags=mode_flags,
                      nvml_process_peak_gib=nvml_peak.stop(), torch_process_peak_gib=max(stage_peaks) if stage_peaks else None,
                      save_artifacts_seconds=round(t_artifacts, 3), save_wav_seconds=round(t_wav, 3),
                      timing=song.timing,
                      audio_seconds=len(song.audio) / song.sample_rate,
                      packed_runtime=pipe.runtime_report(), truncated=song.truncated,
                      allocator_config=os.environ.get("PYTORCH_ALLOC_CONF", os.environ.get("PYTORCH_CUDA_ALLOC_CONF")))
        (out / "release-run.json").write_text(json.dumps(report, indent=2))
        print(json.dumps(report), flush=True)


if __name__ == "__main__":
    main()