"""Whisper transcription — backend-agnostic façade.

Two real backends ship today; the factory picks based on `system_info.cpu_arch()`:

    Apple Silicon  →  MLXTranscriber       (mlx-whisper, Metal-accelerated)
    Intel x86_64   →  WhisperCppTranscriber (pywhispercpp, CPU)

Both expose the same surface so app.py doesn't care which is running:

    t = make_transcriber(cfg)
    t.ensure_model_ready(progress=lambda done, total, msg: ...)
    result = t.transcribe(audio_np, sample_rate=16000)
    # result = {"text": str, "language": str|None, "segments": [...]}

The MLX implementation matches the pre-refactor behaviour byte-for-byte
(lazy load of mlx_whisper, transcribe() takes the same kwargs). The
whisper.cpp implementation is new but produces the same return shape.

Model storage: ~/Library/Application Support/Lumo/models/
"""

from __future__ import annotations

import os
import urllib.request
from pathlib import Path
from typing import Any, Callable, Optional

import numpy as np

from .system_info import (
    BACKEND_MLX,
    BACKEND_WHISPER_CPP,
    is_apple_silicon,
    recommended_backend,
)

# Progress callback signature: (bytes_done, total_bytes, message). Backends
# call this during model download. total may be None when unknown (mlx-whisper
# streams from HF without a Content-Length on every file).
ProgressCb = Callable[[int, Optional[int], str], None]


def model_dir() -> Path:
    """Where downloaded model files live. We use a single root so the
    onboarding window can show a single 'Downloading…' line, and so the
    user can find/delete the cache if storage gets tight."""
    p = Path.home() / "Library" / "Application Support" / "Lumo" / "models"
    p.mkdir(parents=True, exist_ok=True)
    return p


# ── Base ────────────────────────────────────────────────────────────────

class _BaseTranscriber:
    """Backend interface. Subclasses MUST implement the four methods below
    and may override `friendly_name` for the onboarding UI."""

    friendly_name = "whisper"

    def __init__(self, language: str = "auto", initial_prompt: str = "",
                 no_speech_threshold: float = 0.6,
                 condition_on_previous_text: bool = False):
        self.language = None if language == "auto" else language
        self.initial_prompt = initial_prompt or None
        self.no_speech_threshold = no_speech_threshold
        self.condition_on_previous_text = condition_on_previous_text

    def is_model_ready(self) -> bool:
        raise NotImplementedError

    def estimated_size_bytes(self) -> int:
        raise NotImplementedError

    def ensure_model_ready(self, progress: ProgressCb | None = None) -> None:
        """Block until the model is on disk and importable. Idempotent —
        if already ready, returns immediately. Calls `progress` periodically
        during download so the UI can update."""
        raise NotImplementedError

    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> dict[str, Any]:
        raise NotImplementedError


# ── Apple Silicon: mlx-whisper ──────────────────────────────────────────

class MLXTranscriber(_BaseTranscriber):
    """mlx-whisper backend. Model is a Hugging Face repo (NOT a single file)
    so we let huggingface_hub do the download; we just trigger it explicitly
    in ensure_model_ready() instead of waiting for the first transcribe call
    (the legacy behaviour). That keeps the UI responsive on first run."""

    friendly_name = "MLX Whisper (Apple Silicon)"

    # Per HF metadata for whisper-medium-mlx — used to drive the progress bar.
    _APPROX_BYTES = 1_530_000_000  # ~1.5 GB

    def __init__(self, model_repo: str = "mlx-community/whisper-medium-mlx", **kw):
        super().__init__(**kw)
        self.model_repo = model_repo
        self._mlx_whisper = None  # imported lazily — heavy import
        # CR #5 — MLX is NOT thread-safe. Previously app.py held its
        # own _process_lock around dictation transcribes but the meeting
        # live-transcribe worker ran concurrently with no shared lock.
        # Moving the mutex into the transcriber covers every caller so
        # dictation + meeting can't trample each other's MLX state.
        import threading as _th
        self._mlx_lock = _th.Lock()

    def _local_repo_path(self) -> Path | None:
        """Where huggingface_hub cached the snapshot — only counts the model
        as ready if BOTH `config.json` and at least one weights file are on
        disk. The earlier check only required `snapshots/` to be non-empty,
        which was true even when huggingface_hub had just created the
        directory shell on a fresh Mac → the app reported "model ready",
        the onboarding skipped the download page, and the first dictation
        attempt crashed. Returns None on cache miss / partial download."""
        try:
            from huggingface_hub.file_download import repo_folder_name  # type: ignore
            from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE  # type: ignore
            snapshots = Path(HUGGINGFACE_HUB_CACHE) / repo_folder_name(
                repo_id=self.model_repo, repo_type="model"
            ) / "snapshots"
            if not snapshots.exists():
                return None
            # Each subdir under snapshots/ is one revision (usually 1 rev).
            for rev in snapshots.iterdir():
                if not rev.is_dir():
                    continue
                has_config = (rev / "config.json").exists()
                # Whisper-mlx ships either `weights.npz` (older repos) or
                # one+ `.safetensors` shard. Accept either.
                has_weights = (
                    (rev / "weights.npz").exists()
                    or any(rev.glob("*.safetensors"))
                    or any(rev.glob("weights*.npz"))
                )
                if has_config and has_weights:
                    return rev
        except Exception:
            pass
        return None

    def is_model_ready(self) -> bool:
        return self._local_repo_path() is not None

    def estimated_size_bytes(self) -> int:
        return self._APPROX_BYTES

    def ensure_model_ready(self, progress: ProgressCb | None = None) -> None:
        if self.is_model_ready():
            if progress:
                progress(self._APPROX_BYTES, self._APPROX_BYTES, "Model already cached")
            return

        if progress:
            progress(0, self._APPROX_BYTES, f"Downloading {self.model_repo}…")

        # snapshot_download fetches every file in the repo, with parallel
        # workers and a resumable HTTP cache. It's the same path mlx_whisper
        # would have hit lazily; we just call it upfront so the user sees
        # a deterministic 'getting ready' phase.
        from huggingface_hub import snapshot_download
        snapshot_download(repo_id=self.model_repo)

        if progress:
            progress(self._APPROX_BYTES, self._APPROX_BYTES, "Model ready")

    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> dict[str, Any]:
        if audio.size == 0:
            return {"text": "", "language": None, "segments": []}
        if sample_rate != 16000:
            raise ValueError(f"mlx-whisper expects 16kHz, got {sample_rate}")

        if self._mlx_whisper is None:
            import mlx_whisper  # heavy import, lazy
            self._mlx_whisper = mlx_whisper

        kwargs: dict[str, Any] = {
            "path_or_hf_repo": self.model_repo,
            "no_speech_threshold": self.no_speech_threshold,
            "condition_on_previous_text": self.condition_on_previous_text,
            "verbose": None,
        }
        if self.language:
            kwargs["language"] = self.language
        if self.initial_prompt:
            kwargs["initial_prompt"] = self.initial_prompt

        # CR #5 — serialise the actual decode call. Holds for the whole
        # mlx_whisper.transcribe duration so a meeting chunk and a
        # dictation can't overlap inside MLX's internal state.
        with self._mlx_lock:
            result = self._mlx_whisper.transcribe(
                audio.astype(np.float32, copy=False), **kwargs)
        return {
            "text": (result.get("text") or "").strip(),
            "language": result.get("language"),
            "segments": result.get("segments", []),
        }


# ── Intel: whisper.cpp via pywhispercpp ────────────────────────────────

class WhisperCppTranscriber(_BaseTranscriber):
    """whisper.cpp backend for Intel Macs (and as a fallback). Single .bin
    model file, downloaded once into model_dir()."""

    friendly_name = "whisper.cpp (Intel-compatible)"

    # ggml-medium.bin from Hugging Face's ggerganov/whisper.cpp release.
    _MODEL_FILE = "ggml-medium.bin"
    _MODEL_URL = (
        "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-medium.bin"
    )
    _APPROX_BYTES = 1_530_000_000  # ~1.5 GB

    def __init__(self, **kw):
        super().__init__(**kw)
        self._model = None  # pywhispercpp.Model, lazy

    def _model_path(self) -> Path:
        return model_dir() / self._MODEL_FILE

    def is_model_ready(self) -> bool:
        p = self._model_path()
        # Sanity check: file present AND at least 90% of expected size.
        # whisper.cpp segfaults silently on a half-downloaded file, so it's
        # cheaper to re-download than to debug after the fact.
        return p.exists() and p.stat().st_size >= int(self._APPROX_BYTES * 0.9)

    def estimated_size_bytes(self) -> int:
        return self._APPROX_BYTES

    def ensure_model_ready(self, progress: ProgressCb | None = None) -> None:
        path = self._model_path()
        if self.is_model_ready():
            if progress:
                progress(self._APPROX_BYTES, self._APPROX_BYTES, "Model already cached")
            return

        tmp = path.with_suffix(path.suffix + ".part")
        if progress:
            progress(0, self._APPROX_BYTES, f"Downloading {self._MODEL_FILE}…")

        # Stream to a .part file then atomically rename. Avoids leaving a
        # corrupt ggml-medium.bin behind if the download is killed.
        req = urllib.request.Request(
            self._MODEL_URL,
            headers={"User-Agent": "lumo-client/1.0"},
        )
        # `timeout=30` only covers the SOCKET (connect + each read) — not the
        # total wall-clock. A slow-but-progressing connection still completes.
        # A truly stuck socket (TCP black-hole, DNS poisoning, ISP throttling
        # the HF CDN) bubbles up as URLError instead of hanging the onboarding
        # window forever — the caller turns it into a user-visible error.
        with urllib.request.urlopen(req, timeout=30.0) as resp:
            total_hdr = resp.headers.get("Content-Length")
            total = int(total_hdr) if total_hdr else self._APPROX_BYTES
            done = 0
            with open(tmp, "wb") as out:
                while True:
                    chunk = resp.read(1024 * 1024)  # 1 MB
                    if not chunk:
                        break
                    out.write(chunk)
                    done += len(chunk)
                    if progress:
                        progress(done, total, f"Downloading {self._MODEL_FILE}…")
        os.replace(tmp, path)

        if progress:
            progress(self._APPROX_BYTES, self._APPROX_BYTES, "Model ready")

    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> dict[str, Any]:
        if audio.size == 0:
            return {"text": "", "language": None, "segments": []}
        if sample_rate != 16000:
            raise ValueError(f"whisper.cpp expects 16kHz, got {sample_rate}")

        if self._model is None:
            # pywhispercpp ships precompiled wheels for Intel + Apple Silicon.
            # The import is heavy (loads native lib + GGML), do it lazily.
            from pywhispercpp.model import Model  # type: ignore
            self._model = Model(str(self._model_path()))

        # pywhispercpp wants float32 PCM in [-1.0, 1.0]. Our recorder
        # already emits exactly that, so just enforce the dtype.
        segs = self._model.transcribe(
            audio.astype(np.float32, copy=False),
            language=(self.language or "auto"),
            initial_prompt=(self.initial_prompt or ""),
            n_threads=os.cpu_count() or 4,
        )
        text = " ".join(s.text.strip() for s in segs).strip()
        return {
            "text": text,
            "language": getattr(segs[0], "language", None) if segs else None,
            "segments": [
                {"text": s.text, "start": s.t0, "end": s.t1} for s in segs
            ],
        }


# ── Factory ─────────────────────────────────────────────────────────────

def make_dictation_transcriber(full_cfg: dict, fallback_local):
    """Pick the transcriber for the dictation hotkey path specifically.

    Default for v8+: Cloud Mode on → returns a CloudDictator. The user can
    flip `dictation.cloud_mode = false` in config.json to keep dictation
    100% offline (uses the same local transcriber as meetings).

    `fallback_local` is the already-built local transcriber (MLX or
    whisper.cpp) — handed in so app.py can pass cloud errors through it
    when the network / server is unreachable. Done this way (not by
    rebuilding inside) so we don't pay the MLX warmup cost twice."""
    dictation_cfg = full_cfg.get("dictation") or {}
    cloud_on = bool(dictation_cfg.get("cloud_mode", True))
    if not cloud_on:
        return fallback_local

    whisper_cfg = full_cfg.get("whisper") or {}
    try:
        from .cloud_dictator import CloudDictator
        return CloudDictator(
            language=whisper_cfg.get("language", "auto"),
        )
    except Exception as e:  # noqa: BLE001
        # Import error → cloud_dictator module missing from the bundle.
        # Surface in the log; fall back to local so the user still works.
        print(f"[transcriber] cloud dictator unavailable, using local: {e}",
              flush=True)
        return fallback_local


def make_transcriber(cfg: dict) -> _BaseTranscriber:
    """Pick the right LOCAL backend for THIS Mac (used by meetings + by
    dictation when Cloud Mode is off). `cfg` is the parsed config.json
    'whisper' subtree — model name (for mlx), language, prompt, thresholds.
    The user can force a backend via `whisper.backend = "mlx" | "whisper_cpp"`
    in config.json to override auto-detection (useful when MLX is broken on
    a particular OS release and we want to fall back without recompiling)."""
    backend = (cfg.get("backend") or recommended_backend()).lower()
    common = dict(
        language=cfg.get("language", "auto"),
        initial_prompt=cfg.get("initial_prompt", ""),
        no_speech_threshold=cfg.get("no_speech_threshold", 0.6),
        condition_on_previous_text=cfg.get("condition_on_previous_text", False),
    )
    if backend == BACKEND_MLX:
        if not is_apple_silicon():
            print(f"[transcriber] backend=mlx forced on non-arm64 — falling back to whisper.cpp")
            return WhisperCppTranscriber(**common)
        # Defensive: if mlx_whisper isn't importable on this build (e.g. a
        # universal bundle running under Rosetta, or a partial install), fall
        # back to whisper.cpp rather than crash on first dictation. We probe
        # importlib without actually loading the heavy module.
        import importlib.util
        if importlib.util.find_spec("mlx_whisper") is None:
            print("[transcriber] mlx_whisper not importable — falling back to whisper.cpp")
            return WhisperCppTranscriber(**common)
        model_repo = cfg.get("model") or cfg.get("model_mlx") or "mlx-community/whisper-medium-mlx"
        return MLXTranscriber(model_repo=model_repo, **common)
    if backend == BACKEND_WHISPER_CPP:
        return WhisperCppTranscriber(**common)
    raise ValueError(f"unknown whisper.backend: {backend!r}")


# ── Backward compatibility ──────────────────────────────────────────────

class Transcriber(_BaseTranscriber):
    """Legacy alias. Old code that does `Transcriber(model_repo=…)` still
    works — we dispatch to the new MLX class. New code should call
    `make_transcriber(cfg)` instead so the Intel path activates correctly."""

    def __new__(cls, model_repo: str = "mlx-community/whisper-medium-mlx", **kw):
        return MLXTranscriber(model_repo=model_repo, **kw)
