"""Local Dasheng-AudioGen model loading and inference."""

from __future__ import annotations

import os
import threading
import uuid
from pathlib import Path

# Model files preloaded by Spaces can be read-only at runtime. Transformers
# still needs a writable location for code loaded via trust_remote_code=True.
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/huggingface/modules")

try:
    import spaces
except ImportError:  # The package is only required by Hugging Face ZeroGPU.
    spaces = None

import torch
import torchaudio
from transformers import AutoModel


MODEL_ID = os.environ.get(
    "AUDIOGEN_MODEL_ID", "mispeech/Dasheng-AudioGen"
)
DEVICE = os.environ.get("AUDIOGEN_DEVICE", "cuda")
DTYPE_NAME = os.environ.get("AUDIOGEN_DTYPE", "float32").lower()
NUM_STEPS = int(os.environ.get("AUDIOGEN_NUM_STEPS", "25"))
GUIDANCE_SCALE = float(os.environ.get("AUDIOGEN_GUIDANCE_SCALE", "5.0"))
SWAY_SAMPLING_COEF = float(
    os.environ.get("AUDIOGEN_SWAY_SAMPLING_COEF", "-1.0")
)
OUTPUT_DIR = Path(os.environ.get("AUDIOGEN_OUTPUT_DIR", "outputs"))
LOCAL_FILES_ONLY = (
    os.environ.get("AUDIOGEN_LOCAL_FILES_ONLY", "0").lower()
    in {"1", "true", "yes"}
)

_DTYPES = {
    "float32": torch.float32,
    "fp32": torch.float32,
    "float16": torch.float16,
    "fp16": torch.float16,
    "bfloat16": torch.bfloat16,
    "bf16": torch.bfloat16,
}
if DTYPE_NAME not in _DTYPES:
    raise ValueError(
        "AUDIOGEN_DTYPE must be one of: float32, float16, bfloat16."
    )
DTYPE = _DTYPES[DTYPE_NAME]

_INFERENCE_LOCK = threading.Lock()


def _gpu_function(function):
    """Request a ZeroGPU allocation on Spaces and act as a no-op locally."""
    if spaces is None:
        return function
    return spaces.GPU(duration=120)(function)


def _load_model():
    if (
        spaces is None
        and DEVICE.startswith("cuda")
        and not torch.cuda.is_available()
    ):
        raise RuntimeError(
            "Dasheng-AudioGen requires a CUDA GPU, but CUDA is unavailable. "
            "Select GPU or ZeroGPU hardware for the Hugging Face Space."
        )

    model = AutoModel.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        local_files_only=LOCAL_FILES_ONLY,
    )
    model = model.to(device=DEVICE, dtype=DTYPE)
    model.eval()
    return model


# ZeroGPU optimizes model placement performed during application startup.
MODEL = _load_model()


@_gpu_function
def generate_audio_file(structured_prompt: str) -> str:
    """Generate one waveform locally and return its unique WAV file path."""
    prompt = MODEL.compose_prompt(prompt=structured_prompt)

    with _INFERENCE_LOCK, torch.inference_mode():
        audio = MODEL.generate(
            prompt,
            num_steps=NUM_STEPS,
            guidance_scale=GUIDANCE_SCALE,
            sway_sampling_coef=SWAY_SAMPLING_COEF,
        )

    audio = audio.detach().cpu()
    if audio.ndim == 1:
        audio = audio.unsqueeze(0)
    if audio.ndim != 2:
        raise RuntimeError(
            f"Unexpected generated audio shape: {tuple(audio.shape)}"
        )

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    output_path = OUTPUT_DIR / f"audiogen_{uuid.uuid4().hex}.wav"
    sample_rate = int(getattr(MODEL.config, "sample_rate", 16000))
    torchaudio.save(str(output_path), audio, sample_rate)
    return str(output_path)