"""Speech-to-text engine abstraction: local faster-whisper or remote endpoints.

Engines are user-managed presets. Each is either the in-process local
faster-whisper model, or a remote OpenAI-compatible server exposing
`/audio/transcriptions` (faster-whisper-server, Groq, WhisperX, whisper.cpp's
OpenAI shim, NVIDIA NIMs, …). Provides reachability checks (online/offline) and
a benchmark helper (transcript + elapsed seconds).
"""

from __future__ import annotations

import json
import socket
import time
import urllib.error
import urllib.request
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from urllib.parse import urlparse


@dataclass
class STTEngine:
    name: str
    type: str = "local"        # "local" | "openai" | "riva_realtime"
    url: str = ""              # base URL incl. /v1 for remote, e.g. http://localhost:8010/v1
    model: str = ""            # remote model id, or local whisper size override
    api_key_env: str = ""      # env var holding a bearer key (optional)
    timeout: int = 30          # HTTP timeout in seconds; raise for slow servers (WhisperX, etc.)

    @property
    def is_local(self) -> bool:
        return self.type == "local"

    @property
    def is_streaming(self) -> bool:
        return self.type == "riva_realtime"


class STTError(RuntimeError):
    pass


# --- reachability ------------------------------------------------------------
def reachable(url: str, timeout: float = 2.0) -> bool:
    """True if a TCP connection to the URL's host:port succeeds."""
    host, port = _host_port(url)
    if not host:
        return False
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except OSError:
        return False


def status(engine: STTEngine, timeout: float = 2.0) -> bool:
    """True if the engine is usable now (local always; remote = TCP reachable)."""
    if engine.is_local:
        return True
    return reachable(engine.url, timeout)


@dataclass
class ModelMeta:
    """Model id plus optional metadata (languages, etc.) from the server."""
    id: str
    languages: list[str] = field(default_factory=list)


def fmt_languages(langs: list[str]) -> str:
    """Compact display string for a language list, e.g. 'en, de, fr +45'."""
    if not langs:
        return "—"
    if len(langs) >= 50:
        return f"multilingual ({len(langs)})"
    if len(langs) > 5:
        return f"{', '.join(langs[:5])} +{len(langs) - 5}"
    return ", ".join(langs)


def _api_base(url: str) -> str:
    """Strip endpoint-specific path suffixes so metadata/model probes hit the right root.

    http://host:8081/v1/transcribe  →  http://host:8081/v1
    http://host:8081/transcribe     →  http://host:8081
    http://host:8081/v1             →  http://host:8081/v1  (unchanged)
    """
    from urllib.parse import urlparse, urlunparse
    p = urlparse(url.rstrip("/"))
    path = p.path.rstrip("/")
    v1_idx = path.find("/v1")
    if v1_idx >= 0:
        path = path[:v1_idx + 3]   # keep up to and including /v1
    elif path not in ("", "/"):
        path = ""                   # strip unknown custom suffix entirely
    return urlunparse((p.scheme, p.netloc, path, "", "", ""))


def list_models(base_url: str, api_key_env: str = "", timeout: float = 5.0) -> list[str]:
    """Fetch model ids from an OpenAI-compatible, Ollama-style, or Riva/NIM /models endpoint."""
    return [m.id for m in list_models_meta(base_url, api_key_env, timeout)]


def list_models_meta(base_url: str, api_key_env: str = "", timeout: float = 5.0) -> list[ModelMeta]:
    """Like list_models() but returns ModelMeta with language info when available."""
    import os

    if not base_url:
        return []
    base = _api_base(base_url)
    headers: dict[str, str] = {}
    key = os.environ.get(api_key_env) if api_key_env else None
    if key:
        headers["Authorization"] = f"Bearer {key}"

    def _get(url: str) -> dict | None:
        try:
            req = urllib.request.Request(url, headers=headers)
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except (urllib.error.URLError, json.JSONDecodeError, OSError):
            return None

    # 1. Standard OpenAI /models — faster-whisper-server also returns "language"
    data = _get(base + "/models")
    if isinstance(data, dict):
        items = data.get("data")
        if isinstance(items, list):
            result = [ModelMeta(id=m["id"], languages=m.get("language") or [])
                      for m in items if isinstance(m, dict) and m.get("id")]
            if result:
                return result
        items = data.get("models")
        if isinstance(items, list):  # Ollama shape
            return [ModelMeta(id=m.get("name") or m.get("model", ""))
                    for m in items if m.get("name") or m.get("model")]

    # 2. NVIDIA Riva / NIM
    data = _get(base + "/metadata")
    if isinstance(data, dict):
        for info in data.get("modelInfo") or []:
            name = info.get("shortName") or info.get("modelUrl") or ""
            if name:
                return [ModelMeta(id=name.split(":")[0])]

    return []


def probe_server_ram_mb(base_url: str, timeout: float = 3.0) -> float | None:
    """Probe a remote server's Prometheus /metrics endpoint for RSS memory.

    Returns RSS in MB if the server exposes Prometheus metrics with either
    ``process_resident_memory_bytes`` (standard Python/Go exporter) or
    ``container_memory_rss`` (cAdvisor). Returns None if not available.

    Tries the server root first (http://host:port/metrics) then the api base
    path (http://host:port/v1/metrics), because Prometheus endpoints are almost
    always mounted at root even when the API lives under /v1.
    """
    if not base_url:
        return None
    p = urlparse(base_url.rstrip("/"))
    root = f"{p.scheme}://{p.netloc}"
    candidates = [root + "/metrics"]
    api_base_metrics = _api_base(base_url) + "/metrics"
    if api_base_metrics != candidates[0]:
        candidates.append(api_base_metrics)

    def _parse(text: str) -> float | None:
        for line in text.splitlines():
            line = line.strip()
            if line.startswith("#"):
                continue
            for metric in ("process_resident_memory_bytes", "container_memory_rss"):
                if line.startswith(metric):
                    parts = line.split()
                    if len(parts) >= 2:
                        try:
                            return float(parts[-1]) / (1024 * 1024)  # bytes → MB
                        except ValueError:
                            pass
        return None

    for url in candidates:
        try:
            req = urllib.request.Request(url)
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                result = _parse(resp.read().decode("utf-8", "replace"))
                if result is not None:
                    return result
        except Exception:
            continue
    return None


def detect_remote_device(base_url: str, timeout: float = 3.0) -> str:
    """Best-effort GPU/CPU detection for a remote STT server.

    Tries faster-whisper-server's /info endpoint (returns {"device":"cuda",...}),
    then NVIDIA NIM /metadata (GPU-only service). Falls back to "remote".
    """
    if not base_url:
        return "remote"
    base = _api_base(base_url)
    try:
        req = urllib.request.Request(base + "/info")
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        if isinstance(data, dict):
            dev = str(data.get("device") or data.get("compute_type") or "")
            if "cuda" in dev.lower():
                return "CUDA"
            if dev:
                return dev.upper()[:16]
    except Exception:
        pass
    try:
        req = urllib.request.Request(base + "/metadata")
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        if isinstance(data, dict) and data.get("modelInfo"):
            return "CUDA"  # NVIDIA NIM is always GPU
    except Exception:
        pass
    return "remote"


def _host_port(url: str) -> tuple[str | None, int]:
    try:
        u = urlparse(url if "://" in url else "http://" + url)
        port = u.port or (443 if u.scheme == "https" else 80)
        return u.hostname, port
    except ValueError:
        return None, 0


# --- transcription -----------------------------------------------------------
def transcribe(
    engine: STTEngine,
    audio_path: Path,
    *,
    language: str = "",
    hotwords: str = "",
    local_transcriber=None,
    timeout: int = 30,
) -> str:
    if engine.is_local:
        if local_transcriber is None:
            raise STTError("Local engine selected but the model isn't loaded.")
        return local_transcriber.transcribe(audio_path, language=language, hotwords=hotwords)
    if engine.is_streaming:
        raise STTError("Streaming STT engines are live-only. Use a workflow with mode = \"stream\".")
    return _transcribe_remote(engine, audio_path, language=language, prompt=hotwords,
                              timeout=engine.timeout if engine.timeout > 0 else timeout)


def _transcribe_remote(engine: STTEngine, audio_path: Path, *, language: str, prompt: str, timeout: int) -> str:
    import os

    base = engine.url.rstrip("/")
    # If path is non-standard (not empty / /v1 / /v1.0), treat the full URL as
    # the endpoint — supports WhisperX (/transcribe) and other custom paths.
    _path = urlparse(base).path.rstrip("/")
    if _path in ("", "/v1", "/v1.0"):
        endpoint = base + "/audio/transcriptions"
    else:
        endpoint = base
    fields: dict[str, str] = {"response_format": "json"}
    if engine.model:
        fields["model"] = engine.model
    if language:
        fields["language"] = language
    if prompt:
        fields["prompt"] = prompt

    headers = {}
    key = os.environ.get(engine.api_key_env) if engine.api_key_env else None
    if key:
        headers["Authorization"] = f"Bearer {key}"

    body, content_type = _multipart(fields, audio_path)
    headers["Content-Type"] = content_type
    req = urllib.request.Request(endpoint, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:300]
        raise STTError(f"HTTP {exc.code} from {endpoint}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise STTError(f"Cannot reach {endpoint}: {exc.reason}") from exc

    try:
        return (json.loads(raw).get("text") or "").strip()
    except json.JSONDecodeError:
        return raw.strip()  # some servers return plain text


def _multipart(fields: dict, audio_path: Path) -> tuple[bytes, str]:
    boundary = uuid.uuid4().hex
    nl = b"\r\n"
    out = bytearray()
    for k, v in fields.items():
        out += b"--" + boundary.encode() + nl
        out += f'Content-Disposition: form-data; name="{k}"'.encode() + nl + nl
        out += str(v).encode() + nl
    out += b"--" + boundary.encode() + nl
    out += f'Content-Disposition: form-data; name="file"; filename="{audio_path.name}"'.encode() + nl
    out += b"Content-Type: audio/wav" + nl + nl
    out += audio_path.read_bytes() + nl
    out += b"--" + boundary.encode() + b"--" + nl
    return bytes(out), f"multipart/form-data; boundary={boundary}"


# --- benchmark ---------------------------------------------------------------
@dataclass
class BenchResult:
    engine: str
    ok: bool
    text: str = ""
    seconds: float = 0.0
    error: str = ""


def benchmark(engine: STTEngine, audio_path: Path, *, language: str = "", local_transcriber=None) -> BenchResult:
    t0 = time.perf_counter()
    try:
        text = transcribe(engine, audio_path, language=language, local_transcriber=local_transcriber)
        return BenchResult(engine.name, True, text, time.perf_counter() - t0)
    except Exception as exc:  # noqa: BLE001 - report any failure to the UI
        return BenchResult(engine.name, False, "", time.perf_counter() - t0, str(exc))
