# -*- coding: utf-8 -*-
"""
模型预加载工具
在应用启动时预加载所有需要的模型,避免首次请求时的延迟
"""

import logging
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

try:
    from rich.console import Console
except ImportError:
    Console = None

from .boot_events import emit_boot_event

logger = logging.getLogger(__name__)

_PRELOAD_QUIET_LOGGERS = (
    "root",
    "vllm",
    "app.infrastructure.model_utils",
    "app.services.asr.engines.global_models",
    "app.services.asr.qwen3_engine",
    "app.utils.speaker_diarizer",
)


class _ProgressNoiseFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        if record.levelno >= logging.WARNING:
            return True
        return not any(
            record.name == prefix or record.name.startswith(f"{prefix}.")
            for prefix in _PRELOAD_QUIET_LOGGERS
        )


class _StartupProgress:
    def __init__(self, title: str, total: int):
        self._title = title
        self._total = max(total, 1)
        self._enabled = bool(
            Console is not None
            and sys.stderr.isatty()
        )
        self._console: Any = None
        self._filter = _ProgressNoiseFilter()
        self._handlers: list[logging.Handler] = []
        self._current_step = 1
        self._last_description: str | None = None

    def __enter__(self) -> "_StartupProgress":
        emit_boot_event(
            "phase_start",
            phase=self._title,
            total=self._total,
            message=self._title,
        )
        if not self._enabled or Console is None:
            return self
        self._console = Console(stderr=True)
        root_logger = logging.getLogger()
        self._handlers = list(root_logger.handlers)
        for handler in self._handlers:
            handler.addFilter(self._filter)
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        for handler in self._handlers:
            handler.removeFilter(self._filter)
        self._handlers.clear()

    def update(self, description: str) -> None:
        emit_boot_event(
            "step_start",
            phase=self._title,
            step=self._current_step,
            total=self._total,
            message=description,
        )
        if self._console is None:
            return
        if description == self._last_description:
            return
        self._last_description = description
        self._console.print(
            f"[bold cyan][startup {self._current_step}/{self._total}][/bold cyan] {description}",
            highlight=False,
        )

    def advance(self, description: str) -> None:
        emit_boot_event(
            "step_done",
            phase=self._title,
            step=self._current_step,
            total=self._total,
            message=description,
        )
        self._last_description = description
        self._current_step = min(self._current_step + 1, self._total)


@dataclass(frozen=True)
class ModelIntegritySpec:
    description: str
    path: Path
    required_patterns: tuple[str, ...]
    alternative_required_patterns: tuple[tuple[str, ...], ...] = ()
    min_total_size_bytes: int = 0


def _format_bytes(num_bytes: int) -> str:
    value = float(num_bytes)
    units = ["B", "KB", "MB", "GB", "TB"]
    for unit in units:
        if value < 1024.0 or unit == units[-1]:
            return f"{value:.1f}{unit}"
        value /= 1024.0
    return f"{num_bytes}B"


def _find_pattern_matches(root: Path, pattern: str) -> list[Path]:
    return [path for path in root.glob(pattern) if path.is_file()]


def _find_missing_patterns(root: Path, patterns: tuple[str, ...]) -> list[str]:
    return [pattern for pattern in patterns if not _find_pattern_matches(root, pattern)]


def _format_alternative_patterns(pattern_groups: tuple[tuple[str, ...], ...]) -> str:
    return " OR ".join(" + ".join(group) for group in pattern_groups)


def _check_model_integrity_spec(spec: ModelIntegritySpec) -> dict[str, Any]:
    if not spec.path.exists() or not spec.path.is_dir():
        return {
            "description": spec.description,
            "path": str(spec.path),
            "ok": False,
            "missing_patterns": [
                *spec.required_patterns,
                *(
                    [_format_alternative_patterns(spec.alternative_required_patterns)]
                    if spec.alternative_required_patterns
                    else []
                ),
            ],
            "total_size_bytes": 0,
            "reason": "directory_missing",
        }

    files = [path for path in spec.path.rglob("*") if path.is_file()]
    total_size_bytes = sum(path.stat().st_size for path in files)

    missing_patterns = _find_missing_patterns(spec.path, spec.required_patterns)
    if not missing_patterns and spec.alternative_required_patterns:
        alternative_missing_patterns = [
            _find_missing_patterns(spec.path, group)
            for group in spec.alternative_required_patterns
        ]
        if all(alternative_missing_patterns):
            missing_patterns = [
                _format_alternative_patterns(spec.alternative_required_patterns)
            ]

    if missing_patterns:
        return {
            "description": spec.description,
            "path": str(spec.path),
            "ok": False,
            "missing_patterns": missing_patterns,
            "total_size_bytes": total_size_bytes,
            "reason": "required_files_missing",
        }

    if total_size_bytes < spec.min_total_size_bytes:
        return {
            "description": spec.description,
            "path": str(spec.path),
            "ok": False,
            "missing_patterns": [],
            "total_size_bytes": total_size_bytes,
            "reason": "directory_too_small",
        }

    return {
        "description": spec.description,
        "path": str(spec.path),
        "ok": True,
        "missing_patterns": [],
        "total_size_bytes": total_size_bytes,
        "reason": "ok",
    }


def _build_modelscope_spec(
    model_id: str,
    description: str,
    required_patterns: tuple[str, ...],
    *,
    min_total_size_bytes: int,
    alternative_required_patterns: tuple[tuple[str, ...], ...] = (),
) -> ModelIntegritySpec:
    from ..core.config import settings

    return ModelIntegritySpec(
        description=description,
        path=Path(settings.MODELSCOPE_PATH) / model_id,
        required_patterns=required_patterns,
        alternative_required_patterns=alternative_required_patterns,
        min_total_size_bytes=min_total_size_bytes,
    )


def _convert_ms_patterns(
    patterns: tuple[str, ...],
) -> tuple[str, ...]:
    return tuple(p.replace("snapshots/*/", "") for p in patterns)


def _build_qwen_spec(
    model_id: str,
    description: str,
    required_patterns: tuple[str, ...],
    *,
    min_total_size_bytes: int,
    alternative_required_patterns: tuple[tuple[str, ...], ...] = (),
) -> ModelIntegritySpec:
    from ..core.config import settings
    ms_path = Path(settings.MODELSCOPE_PATH) / model_id
    ms_required = _convert_ms_patterns(required_patterns)
    ms_alternative = tuple(
        _convert_ms_patterns(group)
        for group in alternative_required_patterns
    )
    return ModelIntegritySpec(
        description=description,
        path=ms_path,
        required_patterns=ms_required,
        alternative_required_patterns=ms_alternative,
        min_total_size_bytes=min_total_size_bytes,
    )


def _should_check_qwen_forced_aligner(
    resolved_device: str,
    using_cpu_qwen_rust: bool,
) -> bool:
    """Return True when startup integrity should require Qwen forced aligner files."""
    from ..core.config import settings

    _ = (resolved_device, using_cpu_qwen_rust)
    return settings.ASR_ENABLE_WORD_TIMESTAMPS


def _build_required_model_integrity_specs() -> list[ModelIntegritySpec]:
    from ..core.config import settings
    from ..core.device import detect_device
    from ..services.asr.manager import get_model_manager
    from ..services.asr.model_capabilities import (
        get_enabled_qwen_modelscope_assets,
        get_runtime_required_modelscope_assets,
    )
    from ..services.asr.model_plan import get_runtime_model_ids
    from ..services.asr.qwenasr_rust import is_qwenasr_rust_available
    manager = get_model_manager()
    model_ids = [item["id"] for item in manager.list_declared_entries()]
    runtime_models = get_runtime_model_ids(model_ids)
    resolved_device = detect_device(settings.DEVICE)
    using_cpu_qwen_rust = (
        resolved_device == "cpu" and is_qwenasr_rust_available()
    )
    specs: list[ModelIntegritySpec] = []

    for asset in get_runtime_required_modelscope_assets(
        include_realtime_punc=settings.ASR_ENABLE_REALTIME_PUNC,
    ):
        specs.append(
            _build_modelscope_spec(
                asset.model_id,
                asset.description,
                asset.required_patterns,
                alternative_required_patterns=asset.alternative_required_patterns,
                min_total_size_bytes=asset.min_total_size_bytes,
            )
        )

    for asset in get_enabled_qwen_modelscope_assets(
        include_forced_aligner=_should_check_qwen_forced_aligner(
            resolved_device=resolved_device,
            using_cpu_qwen_rust=using_cpu_qwen_rust,
        ),
    ):
        specs.append(
            _build_qwen_spec(
                asset.model_id,
                asset.description,
                asset.required_patterns,
                alternative_required_patterns=asset.alternative_required_patterns,
                min_total_size_bytes=asset.min_total_size_bytes,
            )
        )

    return specs


def verify_required_models_integrity(use_logger: bool = True) -> dict[str, Any]:
    output = logger.info if use_logger else print
    specs = _build_required_model_integrity_specs()
    total = len(specs)
    results: list[dict[str, Any]] = []
    invalid: list[dict[str, Any]] = []

    if not use_logger:
        output("=" * 60)
        output(f"🔍 开始检查运行时模型完整性，共 {total} 个")
        output("=" * 60)
        for index, spec in enumerate(specs, start=1):
            output(f"[{index}/{total}] 检查 {spec.description}")
            result = _check_model_integrity_spec(spec)
            results.append(result)
            if result["ok"]:
                output(
                    f"  ✅ OK  size={_format_bytes(result['total_size_bytes'])} "
                    f"path={result['path']}"
                )
                continue
            invalid.append(result)
            if result["reason"] == "directory_missing":
                output(f"  ❌ FAIL directory_missing path={result['path']}")
            elif result["reason"] == "required_files_missing":
                output(
                    f"  ❌ FAIL missing={', '.join(result['missing_patterns'])} "
                    f"size={_format_bytes(result['total_size_bytes'])} path={result['path']}"
                )
            else:
                output(
                    f"  ❌ FAIL size_too_small size={_format_bytes(result['total_size_bytes'])} "
                    f"path={result['path']}"
                )
        output("=" * 60)
        output(f"模型完整性检查完成: total={total} ok={total - len(invalid)} failed={len(invalid)}")
        output("=" * 60)
        return {
            "total": total,
            "results": results,
            "invalid_models": invalid,
        }

    logger.info("开始检查运行时模型完整性: total=%s", total)
    with _StartupProgress("检查运行时模型完整性", total) as progress:
        for spec in specs:
            progress.update(f"检查 {spec.description}")
            result = _check_model_integrity_spec(spec)
            results.append(result)
            if not result["ok"]:
                invalid.append(result)
                if result["reason"] == "directory_missing":
                    logger.error("模型完整性检查失败: %s, reason=directory_missing, path=%s", spec.description, result["path"])
                elif result["reason"] == "required_files_missing":
                    logger.error(
                        "模型完整性检查失败: %s, reason=required_files_missing, missing=%s, size=%s, path=%s",
                        spec.description,
                        ", ".join(result["missing_patterns"]),
                        _format_bytes(result["total_size_bytes"]),
                        result["path"],
                    )
                else:
                    logger.error(
                        "模型完整性检查失败: %s, reason=directory_too_small, size=%s, path=%s",
                        spec.description,
                        _format_bytes(result["total_size_bytes"]),
                        result["path"],
                    )
            progress.advance(f"检查完成 {spec.description}")

    logger.info(
        "模型完整性检查完成: total=%s ok=%s failed=%s",
        total,
        total - len(invalid),
        len(invalid),
    )

    return {
        "total": total,
        "results": results,
        "invalid_models": invalid,
    }
def preload_models() -> dict[str, Any]:
    """
    预加载所有需要的模型（根据 ENABLE_* 配置过滤）

    Returns:
        dict: 包含加载状态的字典
    """
    # 修复 CAM++ 配置文件（用于离线环境）
    try:
        from .download_models import fix_camplusplus_config
        fix_camplusplus_config()
    except Exception:
        pass  # 修复失败不影响启动

    result: dict[str, Any] = {
        "asr_models": {},  # 所有ASR模型加载状态
        "vad_model": {"loaded": False, "error": None},
        "speaker_diarization_model": {"loaded": False, "error": None},
    }

    from ..core.config import settings
    from ..core.device import detect_device

    # 初始化变量，避免未绑定错误
    asr_device = detect_device(settings.DEVICE)
    model_manager = None

    # 1. 预加载所有配置的ASR模型（根据 ENABLE_* 配置过滤）
    model_ids: list[str] = []
    model_manager = None

    try:
        from ..services.asr.manager import get_model_manager
        from ..services.asr.model_plan import get_runtime_model_ids
        from ..services.asr.runtime import get_runtime_router

        model_manager = get_model_manager()
        runtime_router = get_runtime_router()

        # 获取所有模型配置
        all_models = model_manager.list_declared_entries()
        model_ids = [m["id"] for m in all_models]

        models_to_load = get_runtime_model_ids(model_ids)

        if not models_to_load:
            logger.warning("⚠️  当前环境未解析出可运行的 ASR 模型")

    except Exception as e:
        logger.error(f"❌ 获取模型管理器失败: {e}")
        models_to_load = []
        runtime_router = None

    total_steps = len(models_to_load) + 2

    logger.info(
        "开始预加载模型: declared=%s runtime=%s models=%s",
        len(model_ids) if model_manager else 0,
        len(models_to_load),
        ", ".join(models_to_load) if models_to_load else "（无）",
    )

    with _StartupProgress("预加载模型", total_steps) as progress:
        for model_id in models_to_load:
            result["asr_models"][model_id] = {"loaded": False, "error": None}
            progress.update(f"加载 ASR 模型 {model_id}")
            try:
                if runtime_router is None:
                    raise RuntimeError("runtime router unavailable")
                runtime_router.warmup_model(model_id)
                result["asr_models"][model_id]["loaded"] = True
            except Exception as e:
                result["asr_models"][model_id]["error"] = str(e)
                logger.error("ASR模型预加载失败: %s, error=%s", model_id, e)
            progress.advance(f"已完成 ASR 模型 {model_id}")

        # 2. 预加载语音活动检测模型(VAD)
        progress.update("加载语音活动检测模型(VAD)")
        try:
            from ..services.asr.engines import get_global_vad_model

            vad_model = get_global_vad_model(asr_device)
            if vad_model:
                result["vad_model"]["loaded"] = True
            else:
                result["vad_model"]["error"] = "语音活动检测模型(VAD)加载后返回None"
        except Exception as e:
            result["vad_model"]["error"] = str(e)
            logger.error("语音活动检测模型(VAD)加载失败: %s", e)
        progress.advance("已完成语音活动检测模型(VAD)")

        # 5. 预加载说话人分离模型 (CAM++) - 必需模型，始终加载
        progress.update("加载说话人分离模型(CAM++)")
        try:
            from ..utils.speaker_diarizer import get_global_diarization_pipeline

            diarization_pipeline = get_global_diarization_pipeline()
            if diarization_pipeline:
                result["speaker_diarization_model"]["loaded"] = True
            else:
                result["speaker_diarization_model"]["error"] = "说话人分离模型加载后返回None"
        except Exception as e:
            result["speaker_diarization_model"]["error"] = str(e)
            logger.error("说话人分离模型(CAM++)加载失败: %s", e)
        progress.advance("已完成说话人分离模型(CAM++)")

    loaded_asr_count = sum(1 for status in result["asr_models"].values() if status["loaded"])
    total_asr_count = len(result["asr_models"])
    extra_loaded = sum(
        1
        for key in ("vad_model", "speaker_diarization_model")
        if result[key]["loaded"]
    )
    extra_failed = sum(
        1
        for key in ("vad_model", "speaker_diarization_model")
        if result[key]["error"]
    )
    logger.info(
        "模型预加载完成: asr=%s/%s extra_loaded=%s extra_failed=%s",
        loaded_asr_count,
        total_asr_count,
        extra_loaded,
        extra_failed,
    )

    return result