"""
src/inference.py - Audio8 0.1b inference wrapper (free-tier: ZeroGPU / CPU fallback)
Single model load, BF16 on CUDA, float32 on CPU. 1.7GB total.
"""
import os
import time
import tempfile
from pathlib import Path
from typing import Tuple, Optional

import torch
import soundfile as sf
import librosa

# ZeroGPU support for HF Spaces free tier
try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False
    spaces = None

def _gpu(func):
    """Wrap with spaces.GPU if available (ZeroGPU), else no-op."""
    if HAS_SPACES:
        return spaces.GPU(duration=90)(func)
    return func

MODEL_ID = "Audio8/Audio8-TTS-Preview-0.1b"

# Global singleton
_processor = None
_model = None
_device = None
_dtype = None
_sample_rate = None
_model_loaded = False
_load_error: Optional[str] = None


def get_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def get_dtype(device: str):
    if device == "cuda":
        # bfloat16 halves VRAM vs float32, fits free 16GB
        return torch.bfloat16
    return torch.float32


def is_model_loaded() -> bool:
    return _model_loaded and _model is not None


def get_model_info() -> dict:
    return {
        "model_id": MODEL_ID,
        "device": _device or get_device(),
        "dtype": str(_dtype) if _dtype else str(get_dtype(get_device())),
        "sample_rate": _sample_rate or 44100,
        "loaded": is_model_loaded(),
        "error": _load_error,
        "cuda_available": torch.cuda.is_available(),
    }


def load_model():
    """Load processor + model once. Thread-safe singleton."""
    global _processor, _model, _device, _dtype, _sample_rate, _model_loaded, _load_error
    if is_model_loaded():
        return _processor, _model

    _device = get_device()
    _dtype = get_dtype(_device)

    print(f"[inference] Loading {MODEL_ID} on {_device} ({_dtype}) ...")
    t0 = time.time()

    try:
        from transformers import AutoProcessor, AutoModel

        _processor = AutoProcessor.from_pretrained(
            MODEL_ID, trust_remote_code=True
        )
        _model = AutoModel.from_pretrained(
            MODEL_ID, trust_remote_code=True, dtype=_dtype
        ).eval().to(_device)

        _sample_rate = getattr(_model.config, "codec_sample_rate", 44100)
        _model_loaded = True
        _load_error = None
        print(f"[inference] Loaded in {time.time()-t0:.1f}s | sr={_sample_rate} | device={_device}")
        # Warmup compile check (optional)
        # torch.compile not used on free CPU to avoid high RAM

    except Exception as e:
        _load_error = str(e)
        print(f"[inference] Load failed: {e}")
        raise

    return _processor, _model


def validate_reference_audio(path: str) -> Tuple[bool, str]:
    """Check duration 3-20s, exists, <10MB."""
    if not path or not os.path.exists(path):
        return False, "Reference audio file not found."
    try:
        size_mb = os.path.getsize(path) / (1024 * 1024)
        if size_mb > 10:
            return False, f"File too large ({size_mb:.1f}MB > 10MB)."
        dur = librosa.get_duration(path=path)
        if dur < 1.0:
            return False, f"Audio too short ({dur:.1f}s). Use 3-15s for best clone."
        if dur > 20:
            return False, f"Audio too long ({dur:.1f}s). Trim to <20s."
        if dur < 3:
            # Warning, not blocking
            return True, f"Warning: short audio ({dur:.1f}s) may reduce similarity. 5-15s recommended."
        return True, "OK"
    except Exception as e:
        return False, f"Cannot read audio: {e}"


@_gpu
@torch.inference_mode()
def synthesize(
    reference_audio_path: str,
    reference_text: str,
    text: str,
    temperature: float = 0.7,
    top_p: float = 0.9,
    top_k: int = 50,
    max_new_tokens: int = 512,
) -> Tuple[str, str]:
    """
    Returns (output_wav_path, info_message)
    Raises on failure.
    """
    if not reference_audio_path:
        raise ValueError("Please upload a reference audio (3-15s).")
    if not reference_text or not reference_text.strip():
        raise ValueError("Reference transcript is required and must match the audio exactly.")
    if not text or not text.strip():
        raise ValueError("Please enter text to synthesize.")
    if len(text) > 800:
        raise ValueError(f"Text too long ({len(text)} chars). Limit 800 characters for free tier.")

    ok, msg = validate_reference_audio(reference_audio_path)
    if not ok:
        raise ValueError(msg)

    processor, model = load_model()
    device = _device

    # Processor expects lists
    inputs = processor(
        text=[text.strip()],
        reference_audio=[reference_audio_path],
        reference_text=[reference_text.strip()],
        return_tensors="pt",
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}

    t0 = time.time()
    output = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        do_sample=True,
        return_dict_in_generate=True,
    )
    waveforms, waveform_lengths = model.decode_audio(output.codes)
    gen_time = time.time() - t0

    # Save temp wav
    tmp = tempfile.mktemp(suffix=".wav")
    audio = waveforms[0, : int(waveform_lengths[0])].float().cpu().numpy()
    sf.write(tmp, audio, _sample_rate)

    dur = len(audio) / _sample_rate
    info = f"Generated {dur:.1f}s in {gen_time:.1f}s on {device} | temp={temperature} top_p={top_p} | sr={_sample_rate}"
    if msg.startswith("Warning"):
        info = msg + " | " + info

    return tmp, info


def synthesize_no_clone(
    text: str,
    temperature: float = 0.7,
    top_p: float = 0.9,
    top_k: int = 50,
    max_new_tokens: int = 512,
) -> Tuple[str, str]:
    """TTS without voice cloning (no reference)."""
    if not text or not text.strip():
        raise ValueError("Please enter text.")
    processor, model = load_model()
    device = _device
    inputs = processor(text=[text.strip()], return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.inference_mode():
        out = model.generate(**inputs, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, do_sample=True, return_dict_in_generate=True)
        wavs, lens = model.decode_audio(out.codes)
    tmp = tempfile.mktemp(suffix=".wav")
    audio = wavs[0, : int(lens[0])].float().cpu().numpy()
    sf.write(tmp, audio, _sample_rate)
    return tmp, f"Generated without cloning on {device}"