"""Thin wrapper around Qwen/Qwen3-ForcedAligner-0.6B.

Loads the model once (lazily, on first request) and exposes a single
`align()` function that both the Gradio UI and the Space API call into.
"""

from __future__ import annotations

import logging
import os
import threading
from dataclasses import dataclass

import torch

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("qwen3_force_aligner")

MODEL_ID = os.environ.get("ALIGNER_MODEL_ID", "Qwen/Qwen3-ForcedAligner-0.6B")

# Languages officially supported by Qwen3-ForcedAligner-0.6B, per the model
# card: Chinese, English, Cantonese, French, German, Italian, Japanese,
# Korean, Portuguese, Russian, Spanish. Vietnamese is included below because
# it was requested for this UI; the underlying model may not have been
# trained/evaluated on it, so alignment quality is not guaranteed.
SUPPORTED_LANGUAGES = [
    "Chinese",
    "English",
    "Cantonese",
    "French",
    "German",
    "Italian",
    "Japanese",
    "Korean",
    "Portuguese",
    "Russian",
    "Spanish",
    "Vietnamese",
    "Auto",
]

MAX_AUDIO_SECONDS = 5 * 60  # model card: ~5 minutes per input


@dataclass(frozen=True)
class AlignedSpan:
    index: int
    text: str
    start_time: float
    end_time: float


_model = None
_model_lock = threading.Lock()


def _pick_device_and_dtype() -> tuple[str, torch.dtype]:
    if torch.cuda.is_available():
        return "cuda:0", torch.bfloat16
    if torch.backends.mps.is_available():
        # bf16 forced alignment on MPS is unreliable in practice; fp32 is safe.
        return "mps", torch.float32
    return "cpu", torch.float32


def _flash_attention_available() -> bool:
    try:
        import flash_attn  # noqa: F401

        return True
    except ImportError:
        return False


def get_model():
    """Lazily load and cache the Qwen3ForcedAligner model (thread-safe)."""
    global _model
    if _model is not None:
        return _model

    with _model_lock:
        if _model is not None:
            return _model

        from qwen_asr import Qwen3ForcedAligner

        device, dtype = _pick_device_and_dtype()
        kwargs = dict(dtype=dtype, device_map=device)
        if device.startswith("cuda") and _flash_attention_available():
            kwargs["attn_implementation"] = "flash_attention_2"

        logger.info("Loading %s on %s (%s)...", MODEL_ID, device, dtype)
        _model = Qwen3ForcedAligner.from_pretrained(MODEL_ID, **kwargs)
        logger.info("Model loaded.")
        return _model


def align(audio, text: str, language: str) -> list[AlignedSpan]:
    """Run forced alignment for a single (audio, text, language) triple.

    `audio` may be anything qwen_asr accepts: a local file path, a URL, a
    base64 data URL, or an (np.ndarray, sample_rate) tuple.
    """
    if not audio:
        raise ValueError("An audio input is required.")
    if not text or not text.strip():
        raise ValueError("A text transcript is required.")
    if language not in SUPPORTED_LANGUAGES:
        raise ValueError(f"Unsupported language: {language!r}")

    model = get_model()
    results = model.align(audio=audio, text=text, language=language)
    result = results[0]

    return [
        AlignedSpan(
            index=i,
            text=item.text,
            start_time=float(item.start_time),
            end_time=float(item.end_time),
        )
        for i, item in enumerate(result)
    ]