#!/usr/bin/env python3
"""Run cached autoregressive MuScriptor inference with ONNX Runtime."""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import numpy as np
# Importing PyTorch first preloads the CUDA/cuDNN shared libraries shipped in
# the uv environment when available; CPU-only users do not need PyTorch.
try:
    import torch  # noqa: F401
except ModuleNotFoundError:
    torch = None  # type: ignore[assignment]
import onnxruntime as ort

from muscriptor_onnx.audio import SAMPLE_RATE, load_audio_16k, log_mel_spectrogram


def providers(requested: str, device_id: int) -> list:
    available = ort.get_available_providers()
    if requested == "cuda" or (requested == "auto" and "CUDAExecutionProvider" in available):
        if "CUDAExecutionProvider" not in available:
            raise RuntimeError(f"CUDA EP unavailable; installed providers: {available}")
        return [("CUDAExecutionProvider", {"device_id": device_id}), "CPUExecutionProvider"]
    return ["CPUExecutionProvider"]


def session(path: Path, selected_providers: list) -> ort.InferenceSession:
    options = ort.SessionOptions()
    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    result = ort.InferenceSession(path, sess_options=options, providers=selected_providers)
    requested_cuda = bool(selected_providers) and (
        selected_providers[0] == "CUDAExecutionProvider"
        or (
            isinstance(selected_providers[0], tuple)
            and selected_providers[0][0] == "CUDAExecutionProvider"
        )
    )
    if requested_cuda and "CUDAExecutionProvider" not in result.get_providers():
        raise RuntimeError("CUDA EP was requested but session creation fell back to CPU")
    return result


def causal_mask(query_length: int, past_length: int) -> np.ndarray:
    query_positions = past_length + np.arange(query_length)[:, None]
    key_positions = np.arange(past_length + query_length)[None, :]
    return np.where(key_positions <= query_positions, 0.0, -65504.0).astype(np.float16)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", type=Path, required=True)
    parser.add_argument("--audio", type=Path)
    parser.add_argument("--provider", choices=("auto", "cuda", "cpu"), default="auto")
    parser.add_argument("--device-id", type=int, default=0)
    parser.add_argument("--max-new-tokens", type=int, default=8)
    parser.add_argument(
        "--instrument-id",
        type=int,
        action="append",
        help="Optional MT3_FULL_PLUS group ID; repeat for multiple groups",
    )
    args = parser.parse_args()

    metadata = json.loads((args.model_dir / "config.json").read_text())
    selected = providers(args.provider, args.device_id)
    conditioner = session(args.model_dir / "conditioner.onnx", selected)
    decoder = session(args.model_dir / "decoder.onnx", selected)

    if args.audio:
        audio = load_audio_16k(args.audio)
        source = str(args.audio)
    else:
        # Deterministic smoke-test input; it only tests execution, not quality.
        time_axis = np.arange(5 * SAMPLE_RATE, dtype=np.float32) / SAMPLE_RATE
        audio = (0.1 * np.sin(2 * np.pi * 440.0 * time_axis)).astype(np.float32)
        source = "generated 440 Hz sine"
    mel = log_mel_spectrogram(audio).astype(np.float16)
    instrument_ids = np.asarray(
        [args.instrument_id if args.instrument_id is not None else [-1]], dtype=np.int64
    )
    dataset_ids = np.asarray([[-1]], dtype=np.int64)

    start = time.perf_counter()
    condition = conditioner.run(
        ["condition_embeddings"],
        {
            "log_mel": mel,
            "instrument_ids": instrument_ids,
            "dataset_ids": dataset_ids,
        },
    )[0]
    condition_seconds = time.perf_counter() - start

    layers = metadata["num_layers"]
    heads = metadata["num_heads"]
    head_dim = metadata["head_dim"]
    past_key = np.zeros((layers, 1, heads, 0, head_dim), dtype=np.float16)
    past_value = np.zeros_like(past_key)
    input_ids = np.asarray([[metadata["initial_token_id"]]], dtype=np.int64)
    generated: list[int] = []
    decode_times: list[float] = []

    for step in range(args.max_new_tokens):
        prefix = condition if step == 0 else np.empty((1, 0, heads * head_dim), np.float16)
        query_length = prefix.shape[1] + input_ids.shape[1]
        mask = causal_mask(query_length, past_key.shape[3])
        tick = time.perf_counter()
        logits, past_key, past_value = decoder.run(
            ("logits", "present_key", "present_value"),
            {
                "input_ids": input_ids,
                "condition_embeddings": prefix,
                "past_key": past_key,
                "past_value": past_value,
                "attention_mask": mask,
            },
        )
        decode_times.append(time.perf_counter() - tick)
        if not np.isfinite(logits).all():
            raise RuntimeError("decoder produced non-finite logits")
        logits[:, metadata["first_reserved_token_id"] :] = -np.inf
        token = int(logits.argmax(axis=-1)[0])
        generated.append(token)
        input_ids = np.asarray([[token]], dtype=np.int64)
        if token == metadata["eos_token_id"]:
            break

    active = decoder.get_providers()
    sizes = {
        path.name: path.stat().st_size
        for path in args.model_dir.iterdir()
        if path.is_file() and (path.suffix == ".onnx" or path.name.endswith(".onnx.data"))
    }
    sizes["total"] = sum(sizes.values())
    print(f"model:             {args.model_dir}")
    print(f"audio:             {source}")
    print(f"providers:         {active}")
    print(f"mel shape:         {mel.shape}")
    print(f"condition shape:   {condition.shape}")
    print(f"final KV shape:    {past_key.shape}")
    print(f"generated tokens:  {generated}")
    print(f"condition latency: {condition_seconds * 1000:.1f} ms")
    print(f"decode latency:    {[round(x * 1000, 1) for x in decode_times]} ms")
    print(f"ONNX file sizes:   {sizes}")


if __name__ == "__main__":
    main()