"""
TranscriptionService - faster-whisper based speech-to-text with CUDA/CPU auto-detection.
"""
from __future__ import annotations

import asyncio
import threading
import time
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
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

DEFAULT_MODEL = "medium"
COMPUTE_TYPE_GPU = "float16"
COMPUTE_TYPE_CPU = "int8"
BEAM_SIZE_GPU = 1       # greedy decode — fastest, quality loss negligible on GPU
BEAM_SIZE_CPU = 3       # moderate beam search for CPU (slower but needs quality)
VAD_FILTER = False      # disabled — our real-time VAD already trims the audio
# Maximum audio chunk length for transcription (in samples at 16kHz).
# Whisper was trained on 30s segments.  Longer audio can trigger Windows
# TDR timeout (GPU kernel >2s → process killed silently) on smaller GPUs.
MAX_CHUNK_SAMPLES = 30 * 16000  # 30 seconds


_GPU_VRAM_MB: int | None = None


def gpu_vram_mb() -> int:
    """Total VRAM of the primary NVIDIA GPU in MB (0 = none/undetectable).

    Queried once via nvidia-smi (present whenever the NVIDIA driver is
    installed — no torch import needed) and cached for the process lifetime.
    """
    global _GPU_VRAM_MB
    if _GPU_VRAM_MB is None:
        _GPU_VRAM_MB = _query_gpu_vram_mb()
        logger.info("GPU VRAM detected: {} MB", _GPU_VRAM_MB)
    return _GPU_VRAM_MB


def _nvidia_smi_path() -> str | None:
    """Fully-qualified nvidia-smi path from trusted locations only.

    An unqualified name would be resolved via the process dir/CWD/PATH first —
    the classic Windows binary-planting pattern. Only the driver's fixed
    install locations are accepted.
    """
    import os

    candidates = [
        os.path.join(os.environ.get("SystemRoot", r"C:\Windows"),
                     "System32", "nvidia-smi.exe"),
        os.path.join(os.environ.get("ProgramFiles", r"C:\Program Files"),
                     "NVIDIA Corporation", "NVSMI", "nvidia-smi.exe"),
    ]
    for c in candidates:
        if os.path.isfile(c):
            return c
    return None


def _query_gpu_vram_mb() -> int:
    import os
    import subprocess

    exe = _nvidia_smi_path()
    if exe is None:
        return 0
    try:
        flags = subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0
        proc = subprocess.run(
            [exe, "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=5, creationflags=flags,
        )
        out = proc.stdout.strip()
        if proc.returncode != 0 or not out:
            return 0
        return int(float(out.splitlines()[0].strip()))
    except Exception as e:
        logger.debug("nvidia-smi VRAM query failed: {}", e)
        return 0


_cuda_available_cache: bool | None = None
_cuda_detect_lock = threading.Lock()


def is_cuda_available() -> bool:
    # Cached: the first call initialises the CUDA driver (~1s) while holding
    # the GIL, which visibly freezes the UI. app.run() prewarms this on a
    # worker during startup; later callers (model load, model switch) get
    # the cached answer instantly. The lock makes it single-flight (prewarm
    # and model load can race). Only DETERMINISTIC results are cached —
    # a transient error (driver busy at boot) may succeed on retry.
    global _cuda_available_cache
    if _cuda_available_cache is not None:
        return _cuda_available_cache
    with _cuda_detect_lock:
        if _cuda_available_cache is not None:
            return _cuda_available_cache
        try:
            # Import directly from the C extension to avoid Nuitka compilation
            # issues with ctranslate2's Python subpackages (converters/specs
            # use patterns that break when compiled with Nuitka + Python 3.14).
            from ctranslate2._ext import get_cuda_device_count  # type: ignore[import]
            count = get_cuda_device_count()
            if count > 0:
                logger.info("CUDA available: {} device(s) detected", count)
                _cuda_available_cache = True
                return True
            logger.warning("ctranslate2 loaded but no CUDA devices found (count=0)")
            _cuda_available_cache = False
            return False
        except OSError as e:
            logger.error("CUDA DLL loading failed: {}", e)
            return False
        except Exception as e:
            logger.error("CUDA detection failed: {}", e)
            return False


class TranscriptionService(BaseService):
    def __init__(self, signals: AppSignals, model_name: str = DEFAULT_MODEL) -> None:
        self._signals = signals
        self._model_name = model_name
        self._model = None
        self._ready = False
        self._cuda = False
        self._language: str | None = None
        self._forced_device: str = "auto"  # "auto" | "cuda" | "cpu"
        self._fell_back_to_cpu = False
        # Serialises all model.transcribe calls — the live streaming loop and a
        # possible batch call must never hit the same WhisperModel at once.
        self._model_lock = threading.Lock()

    def set_compute_device(self, device_pref: str) -> None:
        """Set preferred compute device: 'auto', 'cuda', or 'cpu'."""
        self._forced_device = device_pref if device_pref in ("auto", "cuda", "cpu") else "auto"

    def _resolve_device(self) -> tuple[str, str]:
        """Determine actual device and compute type based on preference and availability."""
        hw_cuda = is_cuda_available()

        if self._forced_device == "cpu":
            return "cpu", COMPUTE_TYPE_CPU
        if self._forced_device == "cuda":
            if hw_cuda:
                return "cuda", COMPUTE_TYPE_GPU
            logger.warning("CUDA forced but not available, falling back to CPU")
            return "cpu", COMPUTE_TYPE_CPU
        # auto
        if hw_cuda:
            return "cuda", COMPUTE_TYPE_GPU
        return "cpu", COMPUTE_TYPE_CPU

    async def start(self) -> None:
        model_path = data_dir() / "models" / "whisper" / self._model_name
        if not model_path.exists():
            logger.warning("Whisper model not found at {}", model_path)
            self._signals.error_occurred.emit(
                t("err_model_missing")
            )
            return

        loop = asyncio.get_running_loop()
        # CUDA detection initialises the driver (~1s) — never on the event
        # loop: it froze the waveform of a dictation started during startup.
        device, compute_type = await loop.run_in_executor(None, self._resolve_device)
        self._cuda = device == "cuda"
        self._fell_back_to_cpu = False

        logger.info("Loading Whisper {} on {} ({})", self._model_name, device, compute_type)
        t0 = time.perf_counter()
        try:
            self._model = await loop.run_in_executor(
                None, self._load_model, str(model_path), device, compute_type,
            )
            elapsed = time.perf_counter() - t0
            logger.info("Whisper {} loaded in {:.1f}s on {}", self._model_name, elapsed, device)
            try:
                await self._warmup()
            except Exception as warmup_err:
                logger.error("Whisper warmup failed: {}", warmup_err)
                self._model = None
                raise
            self._ready = True
        except Exception as e:
            logger.error("Failed to load Whisper on {}: {}", device, e)
            if device == "cuda":
                # Automatic CPU fallback
                logger.warning("CUDA failed - falling back to CPU")
                self._cuda = False
                self._fell_back_to_cpu = True
                try:
                    self._model = await loop.run_in_executor(
                        None, self._load_model, str(model_path), "cpu", COMPUTE_TYPE_CPU,
                    )
                    if self._model is None:
                        raise RuntimeError("Model loaded as None after CPU fallback")
                    await self._warmup()
                    self._ready = True
                    self._signals.error_occurred.emit(t("err_gpu_fallback"))
                except Exception as e2:
                    logger.error("CPU fallback also failed: {}", e2)
                    self._model = None
                    self._signals.error_occurred.emit(t("err_speech_start"))
            else:
                self._signals.error_occurred.emit(t("err_speech_start"))

    @property
    def fell_back_to_cpu(self) -> bool:
        return self._fell_back_to_cpu

    def _load_model(self, model_path: str, device: str, compute_type: str) -> object:
        from faster_whisper import WhisperModel  # type: ignore[import]
        return WhisperModel(
            model_size_or_path=model_path,
            device=device,
            compute_type=compute_type,
            cpu_threads=4,
            num_workers=1,
        )

    async def _warmup(self) -> None:
        dummy = np.zeros(16000, dtype=np.float32)
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(None, self._transcribe_sync, dummy)
        logger.debug("Whisper warm-up complete")

    def _transcribe_chunk(self, audio: np.ndarray) -> tuple[str, str | None, float]:
        """Transcribe a single audio chunk (≤30s).

        Returns (text, detected_language, language_probability). The language
        comes from Whisper's detection (or the forced setting with prob 1.0)."""
        beam = BEAM_SIZE_GPU if self._cuda else BEAM_SIZE_CPU
        with self._model_lock:
            segments, info = self._model.transcribe(
                audio,
                language=self._language,
                beam_size=beam,
                best_of=1,
                temperature=0.0,
                condition_on_previous_text=False,
                vad_filter=VAD_FILTER,
                without_timestamps=True,
            )
            text = " ".join(seg.text.strip() for seg in segments).strip()
            lang = getattr(info, "language", None)
            prob = float(getattr(info, "language_probability", 0.0) or 0.0)
            return text, lang, prob

    def transcribe_window_sync(
        self, audio: np.ndarray, language: str | None,
    ) -> list[dict]:
        """Transcribe a short live window with word timestamps.

        Returns [{"raw": word_text_with_spacing, "norm": lowercased_stripped,
        "end": seconds}] for the streaming LocalAgreement loop. Runs under the
        shared model lock so it never overlaps a batch transcription."""
        if self._model is None:
            return []
        beam = BEAM_SIZE_GPU if self._cuda else BEAM_SIZE_CPU
        with self._model_lock:
            segments, _ = self._model.transcribe(
                audio,
                language=language if language else self._language,
                beam_size=beam,
                best_of=1,
                temperature=0.0,
                condition_on_previous_text=False,
                vad_filter=False,          # live audio is already speech-gated
                word_timestamps=True,
            )
            words: list[dict] = []
            for seg in segments:
                for w in (seg.words or []):
                    raw = w.word
                    norm = raw.strip().lower()
                    if norm:
                        words.append({"raw": raw, "norm": norm, "end": float(w.end)})
            return words

    def _transcribe_sync(self, audio: np.ndarray) -> tuple[str, tuple[str, float] | None]:
        """Returns (text, detected_language) where detected_language is
        (code, confidence) or None. Purely call-local — overlapping dictations
        must never see each other's language."""
        if self._model is None:
            return "", None
        # Split long audio into ≤30s chunks to avoid Windows TDR timeout
        # (GPU kernel >2s → OS kills the process silently).
        if len(audio) <= MAX_CHUNK_SAMPLES:
            text, lang, prob = self._transcribe_chunk(audio)
            return text, self._pick_language([(lang, prob)] if text else [])

        parts: list[str] = []
        votes: list[tuple[str | None, float]] = []
        for start in range(0, len(audio), MAX_CHUNK_SAMPLES):
            chunk = audio[start : start + MAX_CHUNK_SAMPLES]
            logger.debug("Transcribing chunk {}-{}s ({} samples)",
                         start // 16000, (start + len(chunk)) // 16000, len(chunk))
            text, lang, prob = self._transcribe_chunk(chunk)
            if text:
                parts.append(text)
                votes.append((lang, prob))
        return " ".join(parts), self._pick_language(votes)

    @staticmethod
    def _pick_language(
        votes: list[tuple[str | None, float]],
    ) -> tuple[str, float] | None:
        """Pick the dictation language from per-chunk votes (probability sum),
        confidence = mean over the chunks that voted for the winner."""
        scores: dict[str, float] = {}
        for lang, prob in votes:
            if lang:
                scores[lang] = scores.get(lang, 0.0) + max(prob, 0.0)
        if not scores:
            return None
        best = max(scores, key=lambda k: scores[k])
        n = sum(1 for lang, _ in votes if lang == best)
        return best, scores[best] / max(n, 1)

    async def transcribe(self, audio: np.ndarray) -> str:
        text, _ = await self.transcribe_with_language(audio)
        return text

    async def transcribe_with_language(
        self, audio: np.ndarray,
    ) -> tuple[str, tuple[str, float] | None]:
        """Transcribe and also return the detected (language, confidence) so
        the AI enhancement can pin the output language. Returned per call, not
        via shared state — overlapping dictations stay isolated."""
        if not self._ready or self._model is None:
            logger.warning("TranscriptionService not ready")
            return "", None
        duration_s = len(audio) / 16000
        logger.info("Transcribing {:.1f}s audio ({} samples)", duration_s, len(audio))
        loop = asyncio.get_running_loop()
        try:
            text, lang = await loop.run_in_executor(None, self._transcribe_sync, audio)
            logger.debug("Transcription: {!r} (lang={})", text[:80] if text else "", lang)
            return text, lang
        except MemoryError:
            logger.error("CUDA OOM on {:.0f}s audio - falling back to CPU", duration_s)
            self._signals.error_occurred.emit(t("err_gpu_oom"))
            return "", None
        except Exception as e:
            logger.error("Transcription error on {:.0f}s audio: {}", duration_s, e)
            self._signals.error_occurred.emit(t("err_transcription"))
            return "", None

    def _transcribe_file_sync(
        self, path: str, language: str | None, with_timestamps: bool,
        on_progress=None, multilingual: bool = False,
    ) -> tuple[str, dict]:
        """Transcribe a media file by path. faster-whisper decodes mp3/ogg/m4a/
        flac/etc. via its bundled PyAV/ffmpeg and windows long audio internally.

        on_progress, if given, is called from this worker thread with a float
        0..1 as decoding advances (per finished segment) — keep it thread-safe
        (Qt signal emits are).
        """
        if self._model is None:
            return "", {}
        beam = BEAM_SIZE_GPU if self._cuda else BEAM_SIZE_CPU
        # In multilingual mode language MUST be None so Whisper auto-detects it
        # (per segment) — a forced language would defeat detection. Meetings
        # may be in a different language than the dictation setting, or switch.
        forced = None if multilingual else (language if language else self._language)
        # Hold the model lock across the whole lazy iteration — segments are
        # produced on demand, so the model is in use until the loop ends.
        with self._model_lock:
            segments, info = self._model.transcribe(
                path,
                language=forced,
                multilingual=multilingual,
                beam_size=beam,
                best_of=1,
                temperature=0.0,
                condition_on_previous_text=False,
                # Files contain real silence (unlike the live-trimmed mic audio),
                # so let Whisper's VAD drop non-speech regions.
                vad_filter=True,
                without_timestamps=not with_timestamps,
            )
            total = float(getattr(info, "duration", 0.0) or 0.0)
            seg_list: list[dict] = []
            parts: list[str] = []
            for seg in segments:
                txt = seg.text.strip()
                if txt:
                    parts.append(txt)
                # Only collect per-segment timing when requested; otherwise
                # start/end are all 0.0 (without_timestamps=True) and waste memory.
                if with_timestamps:
                    seg_list.append({"start": seg.start, "end": seg.end, "text": seg.text})
                if on_progress is not None and total > 0:
                    try:
                        on_progress(min(seg.end / total, 1.0))
                    except Exception:  # progress must never kill the transcription
                        logger.exception("transcription progress callback failed")
        meta = {
            "duration": getattr(info, "duration", None),
            "language": getattr(info, "language", None),
            "segments": seg_list,
        }
        return " ".join(parts).strip(), meta

    async def transcribe_file(
        self,
        path: Path,
        *,
        language: str | None = None,
        with_timestamps: bool = False,
        on_progress=None,
        multilingual: bool = False,
    ) -> tuple[str, dict]:
        """Transcribe an audio/video file. Returns (text, metadata).

        metadata: {"duration": float|None, "language": str|None, "segments": list}
        on_progress: optional callable(float 0..1), invoked from a worker thread.
        multilingual: auto-detect language per segment (for meetings that may be
        in another language than the dictation setting, or switch mid-call).
        """
        if not self._ready or self._model is None:
            logger.warning("TranscriptionService not ready for file transcription")
            return "", {}
        logger.info("Transcribing file: {}", path)
        loop = asyncio.get_running_loop()
        try:
            text, meta = await loop.run_in_executor(
                None, self._transcribe_file_sync, str(path), language,
                with_timestamps, on_progress, multilingual,
            )
            logger.debug("File transcription: {} chars, {:.1f}s audio",
                         len(text), meta.get("duration") or 0.0)
            return text, meta
        except MemoryError:
            logger.error("CUDA OOM transcribing file {}", path)
            self._signals.error_occurred.emit(t("err_gpu_oom"))
            return "", {}
        except Exception as e:
            logger.error("File transcription error for {}: {}", path, e)
            self._signals.error_occurred.emit(t("err_transcription"))
            return "", {}

    async def stop(self) -> None:
        # Refuse new work first, then release the model UNDER the model lock
        # (on a worker: a file transcription may hold it for seconds) — an
        # in-flight job finishes cleanly on the old model instead of hitting
        # a None model mid-transcribe during a hot reload.
        self._ready = False
        loop = asyncio.get_running_loop()

        def _release() -> None:
            with self._model_lock:
                self._model = None

        await loop.run_in_executor(None, _release)

    def set_model(self, model_name: str) -> None:
        from .model_manager import WHISPER_MODELS

        if model_name not in WHISPER_MODELS:
            logger.warning("Invalid model name: {}, keeping {}", model_name, self._model_name)
            return
        self._model_name = model_name

    def set_language(self, lang: str | None) -> None:
        allowed = {
            None, "", "de", "en", "fr", "es", "it", "pl", "pt", "nl", "ru", "ja", "zh",
            "ko", "ar", "tr", "sv", "da", "no", "fi", "cs", "hu", "ro", "uk", "hi",
        }
        if lang not in allowed:
            logger.warning("Invalid language: {}", lang)
            return
        self._language = lang if lang else None

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

    @property
    def model_name(self) -> str:
        return self._model_name

    @property
    def device(self) -> str:
        return "cuda" if self._cuda else "cpu"
