"""Native CUDA runtimes are created inside the allocated ZeroGPU worker."""

import spaces  # Initialize ZeroGPU before importing CUDA-dependent libraries.

import json
import math
import os
import time
from functools import lru_cache
from types import SimpleNamespace

import onnxruntime as ort
import torch
from huggingface_hub import snapshot_download
from transformers import WhisperProcessor

from onnx_runtime import OpenPAOnnx
from asr import transcribe_with_whisperx, prepare_asr_mel
from asr_cache import TranscriptCache

ASR_CACHE = TranscriptCache(max_entries=16, ttl_seconds=600)

def asr_beam_size():
    return max(1, min(5, int(os.getenv("OPENPA_ASR_BEAM_SIZE", "5"))))

if "CUDAExecutionProvider" not in ort.get_available_providers():
    raise RuntimeError("CUDA ONNX Runtime wheel was overwritten by a CPU-only installation")

MODEL_CACHE = os.getenv("HF_HOME", "/tmp/openpa-hf")
MODEL_REVISION = "13a4f4d027c66bf78c36f58f7cf77b06becb67c0"
ASR_REVISION = "edaa852ec7e145841d8ffdb056a99866b5f0a478"


@lru_cache(maxsize=2)
def prepare_artifacts(reference_free):
    """Download on the CPU process, outside GPU quota usage."""
    model = os.getenv("OPENPA_MODEL_DIR") or snapshot_download(
        "fuann/openpa-onnx", revision=MODEL_REVISION,
        token=os.getenv("OPENPA_HF_TOKEN"), cache_dir=MODEL_CACHE,
        allow_patterns=["model.onnx", "model.onnx.data", "*config.json", "tokenizer.json"],
    )
    asr = None
    if reference_free:
        asr = snapshot_download(
            "Systran/faster-whisper-large-v3", revision=ASR_REVISION,
            token=False, cache_dir=MODEL_CACHE,
            allow_patterns=["model.bin", "config.json", "tokenizer.json",
                            "vocabulary.json", "preprocessor_config.json"],
        )
    return model, asr


@lru_cache(maxsize=1)
def get_processor(model_dir):
    return WhisperProcessor.from_pretrained(model_dir, local_files_only=True)


def prepare_request(waveform, rate, text, model_dir, asr_dir):
    """Cache CPU-only processing and imports before allocating a GPU."""
    processor = get_processor(model_dir)
    features = processor.feature_extractor(
        waveform, sampling_rate=rate, return_tensors="np"
    ).input_features.astype("float32")
    if asr_dir:
        import whisperx
        from whisperx.vads.vad import Vad
    return processor, features


def gpu_duration(*args, **kwargs):
    return int(os.getenv("OPENPA_GPU_SECONDS", "90"))


@lru_cache(maxsize=1)
def get_cpu_scorer(model_dir):
    """Keep the CPU session alive in the parent; no CUDA resource is retained."""
    return OpenPAOnnx(model_dir, threads=int(os.getenv("OPENPA_THREADS", "2")),
                      provider="CPUExecutionProvider", processor=get_processor(model_dir))


def asr_duration(waveform, asr_dir, mel_features=None):
    return 12 + math.ceil(len(waveform) / 16000 / 5)


@spaces.GPU(duration=asr_duration)
def transcribe_gpu(waveform, asr_dir, mel_features=None):
    started = time.monotonic()
    if not torch.cuda.is_available():
        raise RuntimeError("ZeroGPU did not allocate a CUDA device")
    options = SimpleNamespace(asr_model=asr_dir, asr_compute_type="float16",
                              asr_batch_size=1, asr_beam_size=asr_beam_size(), local_files_only=True,
                              mel_features=mel_features)
    transcript, compute_type = transcribe_with_whisperx(waveform, options, torch.device("cuda"))
    elapsed = round(time.monotonic() - started, 3)
    print("OPENPA_ASR_GPU_SUCCESS " + json.dumps({"seconds": elapsed,
          "device": torch.cuda.get_device_name(0), "model": "large-v3", "beam_size": asr_beam_size()}), flush=True)
    return transcript, compute_type, elapsed


def run_assessment(waveform, rate, text, model_dir, asr_dir, processor, features, on_stage=None):
    notify = on_stage or (lambda stage: None)
    if os.getenv("OPENPA_SCORING_DEVICE", "cpu") == "cuda":
        notify("transcribe" if text is None else "score")
        return infer_gpu(waveform, rate, text, model_dir, asr_dir, processor, features)
    started = time.monotonic()
    scorer = get_cpu_scorer(model_dir)
    loaded = time.monotonic()
    reference_free = text is None
    asr_seconds = 0
    asr_cache_hit = False
    if reference_free:
        notify("transcribe")
        cache_key = ASR_CACHE.key(waveform, (ASR_REVISION, asr_dir, asr_beam_size(), "float16"))
        cached = ASR_CACHE.get(cache_key)
        if cached is None:
            # WhisperX's log-mel transform runs on CPU. Keep it outside GPU quota.
            mel_features = prepare_asr_mel(waveform)
            text, compute_type, asr_seconds = transcribe_gpu(waveform, asr_dir, mel_features)
            ASR_CACHE.put(cache_key, (text, compute_type))
        else:
            text, compute_type = cached
            asr_cache_hit = True
    transcribed = time.monotonic()
    notify("score")
    inputs, transcript, words = scorer.prepare(waveform, rate, text, input_features=features)
    result = scorer.format_scores(scorer.predict(inputs), transcript, words)
    scored = time.monotonic()
    result["mode"] = "reference-free" if reference_free else "reference-text"
    if reference_free:
        result["asr"] = {"backend": "whisperx", "model": "large-v3",
                         "compute_type": compute_type, "beam_size": asr_beam_size()}
    result["execution"] = {
        "providers": scorer.providers, "asr_cache_hit": asr_cache_hit,
        "timings": {"session_load": round(loaded - started, 3), "asr": asr_seconds,
                    "asr_wall": round(transcribed - loaded, 3),
                    "score": round(scored - transcribed, 3)},
        "compute_seconds": round((loaded - started) + asr_seconds + (scored - transcribed), 3),
        "wall_seconds": round(scored - started, 3),
    }
    print("OPENPA_HYBRID_SUCCESS " + json.dumps({"mode": result["mode"], **result["execution"]}), flush=True)
    return result


@spaces.GPU(duration=gpu_duration)
def infer_gpu(waveform, rate, text, model_dir, asr_dir, processor, features):
    started = time.monotonic()
    if not torch.cuda.is_available():
        raise RuntimeError("ZeroGPU did not allocate a CUDA device")
    ort.preload_dlls()
    ready = time.monotonic()
    scorer = OpenPAOnnx(model_dir, threads=2, provider="CUDAExecutionProvider",
                       processor=processor, cuda_options={"cudnn_conv_algo_search": "HEURISTIC"})
    loaded = time.monotonic()
    reference_free = text is None
    if reference_free:
        options = SimpleNamespace(asr_model=asr_dir, asr_compute_type="float16",
                                  asr_batch_size=1, asr_beam_size=asr_beam_size(), local_files_only=True)
        text, compute_type = transcribe_with_whisperx(waveform, options, torch.device("cuda"))
    transcribed = time.monotonic()
    inputs, transcript, words = scorer.prepare(waveform, rate, text, input_features=features)
    prepared = time.monotonic()
    predictions = scorer.predict(inputs)
    predicted = time.monotonic()
    result = scorer.format_scores(predictions, transcript, words)
    result["mode"] = "reference-free" if reference_free else "reference-text"
    if reference_free:
        result["asr"] = {"backend": "whisperx", "model": "large-v3",
                         "compute_type": compute_type, "beam_size": asr_beam_size()}
    result["execution"] = {
        "providers": scorer.providers, "device": torch.cuda.get_device_name(0),
        "seconds": round(time.monotonic() - started, 3),
        "timings": {"cuda_ready": round(ready - started, 3),
                    "session_load": round(loaded - ready, 3),
                    "asr": round(transcribed - loaded, 3),
                    "token_preparation": round(prepared - transcribed, 3),
                    "forward": round(predicted - prepared, 3)},
        "cuda_options": scorer.session.get_provider_options().get("CUDAExecutionProvider", {}),
    }
    print("OPENPA_GPU_SUCCESS " + json.dumps({
        "mode": result["mode"], **result["execution"]
    }), flush=True)
    return result