"""
ModelManager - Download, verify, and manage Whisper and spell-correction models.
"""
from __future__ import annotations

import asyncio
import hashlib
from pathlib import Path
from typing import TYPE_CHECKING

from loguru import logger

from paths import data_dir
from .base_service import BaseService
from .i18n import t

if TYPE_CHECKING:
    from bridge.signals import AppSignals


# The Systran repos ship vocabulary.txt; newer CT2 conversions (large-v3 family)
# ship vocabulary.json + preprocessor_config.json instead — so the file check
# is defined per model via "files".
_LEGACY_WHISPER_FILES = [
    "model.bin",
    "config.json",
    "tokenizer.json",
    "vocabulary.txt",
]

# "vram_mb" = recommended GPU memory for float16 inference (peak + headroom).
# Below this the model still runs, but on CPU / with OOM risk on GPU.
WHISPER_MODELS: dict[str, dict] = {
    "base": {
        "repo_id": "Systran/faster-whisper-base",
        "dest_rel": "models/whisper/base",
        "size_bytes": 148_000_000,
        "files": _LEGACY_WHISPER_FILES,
        "vram_mb": 700,
    },
    "small": {
        "repo_id": "Systran/faster-whisper-small",
        "dest_rel": "models/whisper/small",
        "size_bytes": 484_000_000,
        "files": _LEGACY_WHISPER_FILES,
        "vram_mb": 1200,
    },
    "medium": {
        "repo_id": "Systran/faster-whisper-medium",
        "dest_rel": "models/whisper/medium",
        "size_bytes": 1_572_000_000,
        "files": _LEGACY_WHISPER_FILES,
        "vram_mb": 2600,
    },
    "large-v3-turbo": {
        "repo_id": "deepdml/faster-whisper-large-v3-turbo-ct2",
        "dest_rel": "models/whisper/large-v3-turbo",
        "size_bytes": 1_631_000_000,
        "files": [
            "model.bin",
            "config.json",
            "preprocessor_config.json",
            "tokenizer.json",
            "vocabulary.json",
        ],
        "vram_mb": 3300,
    },
}


def compat_level(model_name: str, vram_mb: int, device_pref: str = "auto") -> str:
    """How the given model will run on this machine.

    Returns one of: "cpu_forced" (user pinned CPU in settings), "cpu" (no
    usable GPU detected), "gpu_low" (GPU present but below the model's
    recommended VRAM → will end up on CPU / risks OOM), "gpu_ok".
    """
    if device_pref == "cpu":
        return "cpu_forced"
    if vram_mb <= 0:
        return "cpu"
    need = WHISPER_MODELS.get(model_name, {}).get("vram_mb", 0)
    return "gpu_ok" if vram_mb >= need else "gpu_low"


def best_default_model(vram_mb: int) -> str:
    """Best default tier for a fresh install on this hardware.

    Upgrade-only policy: a GPU that comfortably fits the Ultra tier starts
    on it (same quality class as large-v3, faster than medium on GPU);
    everything else keeps the regular "medium" default. Never picks a
    smaller model than the default — baseline quality beats speed.
    """
    if compat_level("large-v3-turbo", vram_mb) == "gpu_ok":
        return "large-v3-turbo"
    return "medium"


SPELL_MODEL = {
    "repo_id": "oliverguhr/spelling-correction-multilingual-base",
    "dest_rel": "models/spell-correction",
    "size_bytes": 900_000_000,
}

REQUIRED_SPELL_FILES = [
    "config.json",
    "tokenizer.json",
    "spiece.model",
]


class ModelManager(BaseService):
    def __init__(self, signals: AppSignals) -> None:
        self._signals = signals
        self._ready = False

    async def start(self) -> None:
        self._ready = True
        logger.info("ModelManager started")

    async def stop(self) -> None:
        self._ready = False

    def whisper_model_exists(self, model_name: str) -> bool:
        if model_name not in WHISPER_MODELS:
            return False
        info = WHISPER_MODELS[model_name]
        dest = data_dir() / info["dest_rel"]
        return all((dest / f).exists() for f in info["files"])

    def spell_model_exists(self) -> bool:
        dest = data_dir() / SPELL_MODEL["dest_rel"]
        return all((dest / f).exists() for f in REQUIRED_SPELL_FILES)

    async def download_whisper(self, model_name: str) -> None:
        if model_name not in WHISPER_MODELS:
            raise ValueError(f"Unknown model: {model_name}")

        info = WHISPER_MODELS[model_name]
        dest = data_dir() / info["dest_rel"]
        dest.mkdir(parents=True, exist_ok=True)

        self._signals.download_status.emit(t("status_downloading_engine"))
        self._signals.download_progress.emit(0.01)
        logger.info("Downloading speech model {} to {}", model_name, dest)

        loop = asyncio.get_running_loop()
        cancel = asyncio.Event()
        monitor = asyncio.create_task(
            self._monitor_dir_progress(dest, info["size_bytes"], cancel)
        )
        try:
            # Prefer getvoxcast's own mirror; fall back to HuggingFace on any error.
            try:
                await loop.run_in_executor(None, self._download_from_voxcast, model_name, str(dest))
            except Exception as e:
                logger.warning(
                    "VoxCast model download failed ({}); falling back to HuggingFace", e
                )
                await loop.run_in_executor(None, self._snapshot_download, info["repo_id"], str(dest))
        finally:
            cancel.set()
            await asyncio.gather(monitor, return_exceptions=True)

        if not self.whisper_model_exists(model_name):
            raise RuntimeError(f"Download completed but model files missing for {model_name}")

        self._signals.download_progress.emit(1.0)
        self._signals.download_status.emit(t("status_engine_ready"))
        logger.info("Speech model {} downloaded", model_name)

    async def download_spell_model(self) -> None:
        dest = data_dir() / SPELL_MODEL["dest_rel"]
        dest.mkdir(parents=True, exist_ok=True)

        self._signals.download_status.emit(t("status_downloading_correction"))
        self._signals.download_progress.emit(0.01)
        logger.info("Downloading spell correction model to {}", dest)

        loop = asyncio.get_running_loop()
        cancel = asyncio.Event()
        monitor = asyncio.create_task(
            self._monitor_dir_progress(dest, SPELL_MODEL["size_bytes"], cancel)
        )
        try:
            await loop.run_in_executor(None, self._snapshot_download, SPELL_MODEL["repo_id"], str(dest))
        finally:
            cancel.set()
            await asyncio.gather(monitor, return_exceptions=True)

        self._signals.download_progress.emit(1.0)
        self._signals.download_status.emit(t("status_correction_ready"))
        logger.info("Spell correction model downloaded")

    # Models hosted on getvoxcast (mirror of the Systran faster-whisper repos).
    MODELS_BASE = "https://getvoxcast.com/downloads/models"

    def _download_from_voxcast(self, model_name: str, local_dir: str) -> None:
        """Download a whisper model's files from getvoxcast and verify SHA-256.
        Raises on any error so the caller can fall back to HuggingFace."""
        import json
        import urllib.request

        hdr = {"User-Agent": "VoxCast"}
        with urllib.request.urlopen(
            urllib.request.Request(f"{self.MODELS_BASE}/manifest.json", headers=hdr), timeout=30
        ) as r:
            manifest = json.load(r)
        model = (manifest.get("models") or {}).get(model_name)
        if not model or not model.get("files"):
            raise RuntimeError(f"No manifest entry for whisper model '{model_name}'")
        dl_base = (manifest.get("base_url") or f"{self.MODELS_BASE}/whisper").rstrip("/")
        out = Path(local_dir)
        out.mkdir(parents=True, exist_ok=True)
        for fname, meta in model["files"].items():
            url = f"{dl_base}/{model_name}/{fname}"
            tmp = out / (fname + ".part")
            with urllib.request.urlopen(
                urllib.request.Request(url, headers=hdr), timeout=180
            ) as resp, open(tmp, "wb") as f:
                while True:
                    chunk = resp.read(1 << 20)
                    if not chunk:
                        break
                    f.write(chunk)
            expected = meta.get("sha256")
            if expected and not self.verify_file_sha256(str(tmp), expected):
                tmp.unlink(missing_ok=True)
                raise RuntimeError(f"SHA-256 mismatch for {fname}")
            tmp.replace(out / fname)
        logger.info("Whisper {} downloaded from getvoxcast", model_name)

    def _snapshot_download(self, repo_id: str, local_dir: str) -> None:
        try:
            from huggingface_hub import snapshot_download  # type: ignore[import]
        except ImportError:
            raise RuntimeError("huggingface_hub not installed")

        snapshot_download(
            repo_id=repo_id,
            local_dir=local_dir,
            local_dir_use_symlinks=False,
        )

    async def _monitor_dir_progress(
        self, dest: Path, expected_bytes: int, cancel: asyncio.Event
    ) -> None:
        """Poll directory size every 0.5s and emit download_progress."""
        _SKIP = {".lock", ".metadata"}
        while not cancel.is_set():
            try:
                total = sum(
                    f.stat().st_size
                    for f in dest.rglob("*")
                    if f.is_file() and f.suffix not in _SKIP
                )
                if expected_bytes > 0:
                    pct = min(0.01 + 0.98 * total / expected_bytes, 0.99)
                    self._signals.download_progress.emit(pct)
            except Exception as e:
                logger.debug("Progress monitor error: {}", e)
            await asyncio.sleep(0.5)

    def verify_model(self, model_name: str) -> bool:
        """Check that all required files exist for a whisper model."""
        return self.whisper_model_exists(model_name)

    def verify_file_sha256(self, path: str, expected: str) -> bool:
        h = hashlib.sha256()
        try:
            with open(path, "rb") as f:
                for chunk in iter(lambda: f.read(65536), b""):
                    h.update(chunk)
            return h.hexdigest() == expected
        except Exception as e:
            logger.error("SHA256 verification failed for {}: {}", path, e)
            return False

    async def verify_all_models(self) -> dict[str, bool]:
        results: dict[str, bool] = {}
        for name in WHISPER_MODELS:
            if self.whisper_model_exists(name):
                results[f"whisper_{name}"] = True
        results["spell_correction"] = self.spell_model_exists()
        return results

    @property
    def is_ready(self) -> bool:
        return self._ready
