"""Local machine-translation engine for #000056 Operation Sandwich.

Two edges, never the middle (ticket §1/§2):

    Spanish query ─▶ [es→en] ─▶ English retrieval + answer + EXISTING
                                 verifier ─▶ [en→es] ─▶ Spanish display

This module only *translates strings*. It has no idea about the
verifier, `audit_mode`, `cache_key`, or `question_hash` — by design.
The caller (`arborist.qa.query`) is responsible for keeping the
grounded core English-only and putting the Spanish rendering in a
display-only field.

The `transformers`/`torch` runtime is the optional `[mt]` extra.
Construction never raises: if the extra is missing or the pinned
checkpoint will not load, :attr:`available` stays False and the
sandwich silently degrades to Phase-0 behaviour (query proceeds
untranslated). This is the `ShadowNLI` pattern (#000049) verbatim.
"""
from __future__ import annotations

import hashlib
import json
import os
import threading
from pathlib import Path
from typing import Optional, Protocol, runtime_checkable

_MANIFEST_PATH = Path(__file__).resolve().parent / "manifest.json"


def load_manifest(path: Optional[Path] = None) -> dict:
    return json.loads((path or _MANIFEST_PATH).read_text())


def _manifest_hash(manifest: dict) -> str:
    """Stable id of the *engine identity* an auditor needs to reproduce
    which sources a translated query retrieved. Pairs (repo+revision) +
    model_version + engine_id only — not the prose `_comment` blobs."""
    core = {
        "mt_model_version": manifest.get("mt_model_version"),
        "engine_id": manifest.get("engine_id"),
        "pairs": {
            k: {"hf_repo": v.get("hf_repo"),
                "pinned_revision": v.get("pinned_revision")}
            for k, v in sorted((manifest.get("pairs") or {}).items())
        },
    }
    canon = json.dumps(core, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(canon.encode("utf-8")).hexdigest()


@runtime_checkable
class Translator(Protocol):
    """Minimal contract the sandwich wiring depends on."""

    available: bool
    engine_id: str
    manifest_hash: str

    def translate(self, text: str, src: str, tgt: str) -> str: ...


class StubTranslator:
    """Deterministic translator — the default test substrate.

    `mapping` is keyed by ``(src, tgt, text)``; an unmapped input is
    returned unchanged so a stub run is a safe identity no-op unless a
    test pins a transform. Deterministic by construction → audit
    replay is byte-stable, like every other stub in the suite.
    """

    def __init__(self, mapping: Optional[dict[tuple[str, str, str], str]] = None,
                 *, fn=None, engine_id: str = "stub-mt", available: bool = True):
        self._map = dict(mapping or {})
        self._fn = fn  # callable(text, src, tgt) -> str; wins over mapping
        self.available = available
        self.engine_id = engine_id
        self.manifest_hash = "stub"

    def translate(self, text: str, src: str, tgt: str) -> str:
        if src == tgt:
            return text
        if self._fn is not None:
            return self._fn(text, src, tgt)
        return self._map.get((src, tgt, text), text)


class OpusMTTranslator:
    """Lazily-loaded pinned Helsinki-NLP opus-mt (MarianMT) pair set.

    Construction never raises (the `ShadowNLI` contract): a missing
    `[mt]` extra or an unloadable checkpoint leaves :attr:`available`
    False and :meth:`translate` returns its input unchanged, so the
    sandwich degrades to Phase-0 (untranslated query) rather than
    erroring.
    """

    def __init__(self, manifest: Optional[dict] = None,
                 device: Optional[str] = None):
        self.manifest = manifest or load_manifest()
        self.engine_id: str = self.manifest.get("engine_id", "opus-mt")
        self.manifest_hash: str = _manifest_hash(self.manifest)
        self.max_length: int = int(self.manifest.get("max_length", 512))
        self.device_pref: str = (
            device or os.environ.get("ARBORIST_MT_DEVICE") or "auto"
        )
        self.device: Optional[str] = None
        self.available = False
        self._reason = "uninitialised"
        self._models: dict[str, tuple] = {}  # "es-en" -> (tok, model)
        # The model load is not concurrency-safe (a parallel bench at
        # concurrency>1 hit `Tensor.item() cannot be called on meta
        # tensors` racing the lazy load). Serialise the *load*; the
        # forward pass is fine concurrently once weights are resident.
        self._load_lock = threading.Lock()

    def _cache_dir(self) -> Path:
        env = os.environ.get("ARBORIST_MT_DIR")
        return Path(env) if env else Path.home() / ".arborist" / "models" / "mt"

    def _ensure_loaded(self) -> None:
        # Double-checked: lock-free fast path once resident/failed,
        # then serialise the one-time load (concurrency-safe).
        if self.available or self._reason.startswith(
            ("deps_missing", "load_failed")
        ):
            return
        with self._load_lock:
            if self.available or self._reason.startswith(
                ("deps_missing", "load_failed")
            ):
                return
            self._load_locked()

    def _load_locked(self) -> None:
        # Deps + device only. Models load lazily per pair (6 manifest
        # pairs × ~300 MB → never load all eagerly; load on first use).
        try:
            import torch  # noqa: F401
            from transformers import (  # noqa: F401
                AutoModelForSeq2SeqLM, AutoTokenizer,
            )
        except ImportError as e:
            self._reason = (
                f"deps_missing: {e} (install: pip install 'arborist[mt]')"
            )
            return
        if self.device_pref == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = self.device_pref
        self.available = True  # engine usable; pairs load on demand
        self._reason = "ok"

    def _ensure_pair(self, pair: str):
        """Load one es-en/en-fr/… pair on first use, serialised (the
        load is the concurrency-unsafe step — #000056 fan-out)."""
        got = self._models.get(pair)
        if got is not None:
            return got
        spec = (self.manifest.get("pairs") or {}).get(pair)
        if spec is None:
            return None
        with self._load_lock:
            got = self._models.get(pair)
            if got is not None:
                return got
            from transformers import (
                AutoModelForSeq2SeqLM, AutoTokenizer,
            )
            try:
                cache = str(self._cache_dir())
                tok = AutoTokenizer.from_pretrained(
                    spec["hf_repo"], revision=spec.get("pinned_revision"),
                    cache_dir=cache,
                )
                model = AutoModelForSeq2SeqLM.from_pretrained(
                    spec["hf_repo"], revision=spec.get("pinned_revision"),
                    cache_dir=cache,
                )
                model.eval()
                if self.device == "cuda":
                    model = model.to("cuda")
            except Exception as e:  # noqa: BLE001 — load failure = passthrough
                self._reason = f"load_failed[{pair}]: {type(e).__name__}: {e}"
                return None
            self._models[pair] = (tok, model)
            return self._models[pair]

    def translate(self, text: str, src: str, tgt: str) -> str:
        if not text or src == tgt:
            return text
        self._ensure_loaded()
        if not self.available:
            return text  # graceful degrade → sandwich no-ops to Phase-0
        got = self._ensure_pair(f"{src}-{tgt}")
        if got is None:
            return text  # unsupported / unloadable pair → safe passthrough
        import torch

        tok, model = got
        enc = tok(
            [text], return_tensors="pt", truncation=True,
            max_length=self.max_length,
        )
        if self.device == "cuda":
            enc = {k: v.to("cuda") for k, v in enc.items()}
        with torch.no_grad():
            out = model.generate(**enc, max_length=self.max_length)
        return tok.batch_decode(out, skip_special_tokens=True)[0].strip() or text


_SINGLETON: Optional[OpusMTTranslator] = None
_SINGLETON_LOCK = threading.Lock()


def get_translator(manifest: Optional[dict] = None) -> Translator:
    """Factory the sandwich uses. Returns a **process-level memoised**
    `OpusMTTranslator` so the ~300 MB MarianMT pair loads **once** and
    is reused across every query (a bench at concurrency>1 reloaded it
    per call → meta-tensor races + 75× waste; #000056 fan-out). The
    caller gates on `.available` (False when `[mt]` is absent → the
    sandwich degrades to Phase-0, never raises). A custom `manifest`
    bypasses the cache (test isolation)."""
    if manifest is not None:
        return OpusMTTranslator(manifest=manifest)
    global _SINGLETON
    if _SINGLETON is None:
        with _SINGLETON_LOCK:
            if _SINGLETON is None:
                _SINGLETON = OpusMTTranslator()
    return _SINGLETON
