"""模型管理器：ASR 与翻译模型不共驻，任一时刻 GPU 上只有一个模型。

设计：
- 单例 ModelManager 跟踪当前加载的模型类型（none / asr / translator）
- get_asr(): 若翻译器在内存 → 先卸载（del + gc + empty_cache）→ 加载 faster-whisper
- get_translator(): 若 ASR 在内存 → 先卸载 → 加载 NLLB
- 翻译阶段独占显存，可用大 batch_size；ASR 阶段同理

这样在 24G 3090 上无需担心显存叠加，CPU dev 时也省内存。
"""

from __future__ import annotations

import gc
import logging
import threading
from typing import Any

from ..config import get_settings

logger = logging.getLogger("audio2text.models")

# 全局单例 + 锁：模型加载/卸载必须串行
_lock = threading.Lock()


def _free_memory() -> None:
    """释放 Python 对象与 GPU 缓存。"""
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
    except Exception:  # pragma: no cover
        pass


class ModelManager:
    """ASR / 翻译模型的不共驻管理器。"""

    def __init__(self) -> None:
        self._asr: Any = None       # faster_whisper.WhisperModel
        self._translator: Any = None  # transformers pipeline
        self._current: str = "none"  # none | asr | translator

    # ---------------- ASR ----------------

    def get_asr(self) -> Any:
        """返回已加载的 faster-whisper 模型；必要时先卸载翻译器。"""
        with _lock:
            if self._asr is not None:
                return self._asr
            if self._translator is not None:
                self._unload_translator_locked()
            s = get_settings().asr
            logger.info("加载 ASR 模型 model=%s device=%s compute_type=%s",
                        s.model, s.device, s.compute_type)
            from faster_whisper import WhisperModel, BatchedInferencePipeline
            # device/compute_type 组合：cpu+int8 / cuda+float16
            # BatchedInferencePipeline 包装 WhisperModel，使 transcribe() 支持 batch_size，
            # 多个音频块（chunk_length=30s）一次性送 GPU 解码，配合内部 prefill 提高利用率。
            whisper = WhisperModel(
                s.model, device=s.device, compute_type=s.compute_type,
            )
            self._asr = BatchedInferencePipeline(model=whisper)
            self._current = "asr"
            logger.info("ASR 模型已就绪（batched, batch_size=%d, beam_size=%d）。",
                        s.batch_size, s.beam_size)
            return self._asr

    def unload_asr(self) -> None:
        with _lock:
            self._unload_asr_locked()

    def _unload_asr_locked(self) -> None:
        if self._asr is None:
            return
        logger.info("卸载 ASR 模型（释放显存供翻译器独占）。")
        # faster-whisper 模型无显式 close，del 即可
        del self._asr
        self._asr = None
        self._current = "none" if self._translator is None else "translator"
        _free_memory()

    # ---------------- 翻译器 ----------------

    def get_translator(self) -> Any:
        """返回已加载的 NLLB 翻译 pipeline；必要时先卸载 ASR。"""
        with _lock:
            if self._translator is not None:
                return self._translator
            if self._asr is not None:
                self._unload_asr_locked()
            s = get_settings().translation
            logger.info("加载翻译模型 model=%s device=%s", s.model, s.device)
            from transformers import pipeline
            self._translator = pipeline(
                "translation",
                model=s.model,
                device=s.device,
                src_lang=s.src_lang,
                tgt_lang=s.tgt_lang,
                batch_size=s.batch_size,  # pipeline 内部批大小，与 translate_service 分块对齐
            )
            self._current = "translator"
            logger.info("翻译模型已就绪（显存独占，batch_size=%d，sort_by_length=%s）。",
                        s.batch_size, s.sort_by_length)
            return self._translator

    def unload_translator(self) -> None:
        with _lock:
            self._unload_translator_locked()

    def _unload_translator_locked(self) -> None:
        if self._translator is None:
            return
        logger.info("卸载翻译模型。")
        # 释放 pipeline 持有的 model + tokenizer
        mdl = getattr(self._translator, "model", None)
        tok = getattr(self._translator, "tokenizer", None)
        del self._translator, mdl, tok
        self._translator = None
        self._current = "none" if self._asr is None else "asr"
        _free_memory()

    # ---------------- 状态 ----------------

    @property
    def current(self) -> str:
        return self._current


# 进程级单例
_manager: ModelManager | None = None


def get_model_manager() -> ModelManager:
    global _manager
    if _manager is None:
        _manager = ModelManager()
    return _manager