from __future__ import annotations

import io
import os
from functools import lru_cache

import numpy as np

MODEL_ID = "CohereLabs/cohere-transcribe-03-2026"


@lru_cache(maxsize=1)
def load_model():
    import torch
    from transformers import AutoProcessor, CohereAsrForConditionalGeneration

    token = os.environ.get("HF_TOKEN")
    if not token:
        raise RuntimeError(
            "HF_TOKEN is required. Set the third-eye-hf Modal secret with a valid token "
            "that has access to CohereLabs/cohere-transcribe-03-2026."
        )

    # PERF FIX: load in half precision and pin to a single GPU explicitly.
    # The previous `device_map="auto"` + default float32 caused the model to run
    # ~1s/forward-pass (effectively on CPU / partially offloaded), making each
    # transcription take minutes. bf16 on a single CUDA device matches the
    # model card's "blazing fast" benchmark (RTFx ~500).
    cuda = torch.cuda.is_available()
    dtype = (
        torch.bfloat16
        if cuda and torch.cuda.is_bf16_supported()
        else (torch.float16 if cuda else torch.float32)
    )

    try:
        processor = AutoProcessor.from_pretrained(
            MODEL_ID,
            token=token,
        )
        model = CohereAsrForConditionalGeneration.from_pretrained(
            MODEL_ID,
            token=token,
            torch_dtype=dtype,
        )
    except OSError as e:
        if "gated" in str(e).lower() or "403" in str(e):
            raise RuntimeError(
                f"HF token does not have access to {MODEL_ID}. "
                f"Visit https://huggingface.co/{MODEL_ID} and accept the license agreement, "
                f"then update the third-eye-hf Modal secret."
            ) from e
        raise

    if cuda:
        model = model.to("cuda")
    model.eval()

    # Diagnostic: proves on the next deploy whether the model is truly on GPU.
    # If `device` is cpu here, that is the real cause of slow transcription.
    param = next(model.parameters())
    print(
        f"[third-eye STT] loaded {MODEL_ID} | cuda_available={cuda} "
        f"| device={param.device} | dtype={param.dtype}",
        flush=True,
    )
    return processor, model


def transcribe_wav_bytes(audio_bytes: bytes, language: str = "en") -> str:
    import time

    import librosa
    import soundfile as sf
    import torch

    audio, sample_rate = sf.read(io.BytesIO(audio_bytes), dtype="float32")
    if audio.ndim > 1:
        audio = np.mean(audio, axis=1)
    if sample_rate != 16_000:
        audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=16_000)

    # ACCURACY FIX: strip leading/trailing silence. The ASR model hallucinates
    # plausible-but-fake words over silent padding (e.g. appending "Yes, sir."
    # after the real question), which then corrupts the question sent to the
    # vision model. Trimming silence removes the trigger. Fall back to the raw
    # audio if trimming would leave less than ~0.1s of signal.
    trimmed, _ = librosa.effects.trim(audio, top_db=30)
    if trimmed.size >= 1_600:
        audio = trimmed

    processor, model = load_model()
    inputs = processor(
        audio,
        sampling_rate=16_000,
        return_tensors="pt",
        language=language,
    )
    inputs = inputs.to(model.device, dtype=model.dtype)
    # PERF FIX: force greedy decoding (num_beams=1, do_sample=False). If the model's
    # generation_config defaulted to beam search, every token cost several forward
    # passes (the repeated "6/6" progress bars in the logs). Greedy is also the most
    # faithful choice for transcription. max_new_tokens trimmed to 128 (speech is short).
    start = time.time()
    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=128,
            num_beams=1,
            do_sample=False,
        )
    print(
        f"[third-eye STT] generate: {time.time() - start:.2f}s "
        f"for {int(generated_ids.shape[-1])} tokens",
        flush=True,
    )
    decoded = processor.decode(generated_ids, skip_special_tokens=True)
    if isinstance(decoded, list):
        decoded = decoded[0]
    return str(decoded).strip()