"""Ultimo's translation service: Mozilla's Firefox Translations models (the
Bergamot engine) and fastText language identification, behind a small HTTP
API that only the webmail API calls (backend/app/translation.py), on the
internal network. A message's text never leaves the server.

- Firefox's models translate to and from English; any other pair goes
  through English (Polish → English → German). So a language is installed as
  its two models, to and from English, and English is always there.
- What is installed is kept in ``$DATA_DIR/state.json``, with the languages
  still to install. The first start asks for ``TRANSLATOR_LANGUAGES``; after
  that the admin panel adds and removes languages (``POST``/``DELETE
  /languages/{code}``). A language that could not be downloaded (no network
  yet) is tried again in the background, at growing intervals, until it is
  in — across restarts too.
- Models come from Mozilla's Remote Settings (where Firefox gets them), each
  file checked against its SHA-256; only release models (no test builds).
- Language identification: fastText's lid.176 (176 languages), baked into
  the image.

API (JSON):
  GET    /health
  GET    /languages            installed / installing / available, errors
  POST   /languages/{code}     install (in the background)
  DELETE /languages/{code}     remove
  POST   /detect     {"q": str | [str]}                     → language(s)
  POST   /translate  {"q": str | [str], "source", "target"} → translation(s)
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import re
import shutil
import threading
import time
import urllib.request
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

import fasttext
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel


def _load_engine() -> Any:
    """The Bergamot engine: only its compiled module. Importing the bergamot
    package itself fetches model catalogues from translatelocally.com and
    OPUS on every start (and fails without them); we load our own models."""
    import importlib.machinery
    import importlib.util

    spec = importlib.util.find_spec("bergamot")
    if spec is None or not spec.submodule_search_locations:
        raise ImportError("the bergamot package is not installed")
    folder = Path(next(iter(spec.submodule_search_locations)))
    path = next(folder.glob("_bergamot*.so"))
    loader = importlib.machinery.ExtensionFileLoader("_bergamot", str(path))
    module_spec = importlib.util.spec_from_loader("_bergamot", loader)
    assert module_spec is not None
    module = importlib.util.module_from_spec(module_spec)
    loader.exec_module(module)
    return module


bergamot = _load_engine()

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("translator")

DATA = Path(os.environ.get("DATA_DIR", "/data"))
MODELS = DATA / "models"
STATE = DATA / "state.json"
LID_MODEL = os.environ.get("LID_MODEL", "/opt/lid.176.ftz")
DEFAULT_LANGUAGES = [
    c.strip() for c in os.environ.get("TRANSLATOR_LANGUAGES", "pl,de,fr,es,it").split(",")
]
THREADS = int(os.environ.get("TRANSLATOR_THREADS", "2"))
RECORDS_URL = os.environ.get(
    "MODELS_RECORDS_URL",
    "https://firefox.settings.services.mozilla.com/v1/buckets/main/collections/"
    "translations-models/records",
)
ATTACHMENTS_URL = os.environ.get(
    "MODELS_ATTACHMENTS_URL", "https://firefox-settings-attachments.cdn.mozilla.net/"
)
PIVOT = "en"
_RECORDS_TTL = 3600
_CODE = re.compile(r"^[a-z]{2,3}(-[A-Za-z]{2,4})?$")

_service = bergamot.Service(bergamot.ServiceConfig(numWorkers=THREADS, logLevel="off"))
_lid = fasttext.load_model(LID_MODEL)
_lock = threading.Lock()  # state, loaded models, installs
_loaded: dict[str, Any] = {}  # "en-pl" → bergamot model
_installing: dict[str, dict[str, int]] = {}  # "pl" → {"done": bytes, "total": bytes}
_errors: dict[str, str] = {}  # "pl" → why the last install failed
_records: tuple[float, list[dict[str, Any]]] = (0.0, [])


# --- what is installed --------------------------------------------------------------


def _read_state() -> dict[str, Any] | None:
    try:
        return json.loads(STATE.read_text())
    except (OSError, ValueError):
        return None


def _installed() -> list[str]:
    state = _read_state() or {}
    return [PIVOT] + [c for c in state.get("languages", []) if c != PIVOT]


def _pending() -> list[str]:
    """Languages asked for and not installed yet (tried again until they are)."""
    return list((_read_state() or {}).get("pending", []))


def _write_state(installed: list[str], pending: list[str]) -> None:
    DATA.mkdir(parents=True, exist_ok=True)
    tmp = STATE.with_suffix(".tmp")
    done = set(installed) - {PIVOT}
    tmp.write_text(
        json.dumps(
            {
                "languages": sorted(done),
                "pending": [c for c in dict.fromkeys(pending) if c not in done and c != PIVOT],
            }
        )
    )
    tmp.replace(STATE)


def _write_installed(codes: list[str]) -> None:
    _write_state(codes, _pending())


# --- Mozilla's model catalogue -----------------------------------------------------------


def _fetch(url: str, timeout: float = 60) -> bytes:
    req = urllib.request.Request(url, headers={"User-Agent": "ultimo-translator"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.read()


def _catalogue() -> list[dict[str, Any]]:
    """Mozilla's model records (cached for an hour)."""
    global _records
    stamp, cached = _records
    if cached and time.monotonic() - stamp < _RECORDS_TTL:
        return cached
    data = json.loads(_fetch(RECORDS_URL))["data"]
    _records = (time.monotonic(), data)
    return data


def _version(text: str) -> tuple[int, ...] | None:
    """ "2.1" → (2, 1); test builds ("2.1a1") → None."""
    return tuple(int(x) for x in text.split(".")) if re.fullmatch(r"\d+(\.\d+)*", text) else None


def _release_files(source: str, target: str) -> list[dict[str, Any]]:
    """The newest release model of one direction: its files' records."""
    by_version: dict[tuple[int, ...], list[dict[str, Any]]] = {}
    for r in _catalogue():
        if r.get("fromLang") != source or r.get("toLang") != target:
            continue
        if r.get("filter_expression"):
            continue  # a test build (Nightly only)
        v = _version(str(r.get("version", "")))
        if v is not None:
            by_version.setdefault(v, []).append(r)
    if not by_version:
        raise LookupError(f"no model for {source} → {target}")
    return by_version[max(by_version)]


def _available() -> list[str]:
    """Languages Mozilla has release models for, both ways with English."""
    pairs: set[tuple[str, str]] = set()
    for r in _catalogue():
        if not r.get("filter_expression") and _version(str(r.get("version", ""))):
            pairs.add((r["fromLang"], r["toLang"]))
    return sorted({a for a, b in pairs if b == PIVOT and (PIVOT, a) in pairs} | {PIVOT})


# --- installing ---------------------------------------------------------------------------


def _model_dir(source: str, target: str) -> Path:
    return MODELS / f"{source}-{target}"


def _install_direction(code: str, source: str, target: str) -> None:
    files = _release_files(source, target)
    final = _model_dir(source, target)
    tmp = final.with_name(final.name + ".partial")
    shutil.rmtree(tmp, ignore_errors=True)
    tmp.mkdir(parents=True)
    for r in files:
        att = r["attachment"]
        blob = _fetch(ATTACHMENTS_URL + att["location"], timeout=600)
        if hashlib.sha256(blob).hexdigest() != att["hash"]:
            raise ValueError(f"{att['filename']}: checksum mismatch")
        (tmp / att["filename"]).write_bytes(blob)
        with _lock:
            _installing[code]["done"] += len(blob)
    (tmp / "files.json").write_text(
        json.dumps(
            {r["fileType"]: r["attachment"]["filename"] for r in files}
            | {"version": files[0]["version"]}
        )
    )
    shutil.rmtree(final, ignore_errors=True)
    tmp.replace(final)


def _install(code: str) -> None:
    """Both directions of ``code`` with English, then mark it installed."""
    try:
        directions = [(code, PIVOT), (PIVOT, code)]
        total = sum(r["attachment"]["size"] for s, t in directions for r in _release_files(s, t))
        with _lock:
            _installing[code]["total"] = total
        for source, target in directions:
            _install_direction(code, source, target)
        with _lock:
            _write_state(_installed() + [code], [c for c in _pending() if c != code])
            for pair in (f"{code}-{PIVOT}", f"{PIVOT}-{code}"):
                _loaded.pop(pair, None)  # a reinstall: load the new files
            _errors.pop(code, None)
        log.info("installed %s", code)
    except Exception as exc:  # noqa: BLE001 — reported in /languages, tried again later
        log.warning("installing %s failed: %s", code, exc)
        with _lock:
            _errors[code] = str(exc)[:300]
    finally:
        with _lock:
            _installing.pop(code, None)


_wake = threading.Event()  # a new language to install: try now, not after the pause
_RETRY_FIRST, _RETRY_MAX = 30.0, 900.0


def _installer() -> None:
    """Install the pending languages, one after another; while some fail
    (no network, Mozilla out of reach), try again after a pause that grows
    from half a minute to a quarter of an hour."""
    pause = _RETRY_FIRST
    while True:
        for code in _pending():
            with _lock:
                if code in _installing or code in _installed():
                    continue
                _installing[code] = {"done": 0, "total": 0}
            _install(code)
        if not _pending():
            pause = _RETRY_FIRST
            _wake.wait()  # nothing to do until a language is asked for
        else:
            _wake.wait(pause)
            pause = min(pause * 2, _RETRY_MAX)
        _wake.clear()


def _ask_for(code: str) -> None:
    with _lock:
        _write_state(_installed(), _pending() + [code])
        _errors.pop(code, None)
    _wake.set()


def _start() -> None:
    """The first start asks for the default languages; every start resumes
    what is still pending."""
    state = _read_state()
    # No state yet — or one from before pending languages were kept, with
    # nothing installed (a first start that could not download): ask for the
    # defaults.
    if state is None or ("pending" not in state and not state.get("languages")):
        _write_state([], [c for c in DEFAULT_LANGUAGES if c])
    threading.Thread(target=_installer, daemon=True).start()


@asynccontextmanager
async def _lifespan(_: FastAPI) -> AsyncIterator[None]:
    _start()
    yield


app = FastAPI(title="Ultimo translator", docs_url=None, redoc_url=None, lifespan=_lifespan)


# --- translating -------------------------------------------------------------------------


def _config(directory: Path) -> str:
    files = json.loads((directory / "files.json").read_text())
    model = files["model"]
    vocabs = (
        [files["vocab"], files["vocab"]]
        if "vocab" in files
        else [files["srcvocab"], files["trgvocab"]]
    )
    # The models are 8-bit; their file name says which kind (as Firefox decides).
    precision = "int8shiftAlphaAll" if model.endswith("alphas.bin") else "int8shiftAll"
    lines = [
        "relative-paths: false",
        "models:",
        f"  - {directory / model}",
        "vocabs:",
        *[f"  - {directory / v}" for v in vocabs],
        "beam-size: 1",
        "normalize: 1.0",
        "word-penalty: 0",
        "mini-batch: 64",
        "maxi-batch: 1000",
        "maxi-batch-sort: src",
        "workspace: 2000",
        "max-length-factor: 2.5",
        "max-length-break: 128",
        "mini-batch-words: 1024",
        "skip-cost: true",
        "quiet: true",
        f"gemm-precision: {precision}",
    ]
    if "lex" in files:
        lines += ["shortlist:", f"  - {directory / files['lex']}", "  - false"]
    return "\n".join(lines) + "\n"


def _model(source: str, target: str) -> Any:
    pair = f"{source}-{target}"
    with _lock:
        if pair not in _loaded:
            directory = _model_dir(source, target)
            if not (directory / "files.json").exists():
                raise HTTPException(400, f"{source} → {target} is not installed")
            _loaded[pair] = _service.modelFromConfig(_config(directory))
        return _loaded[pair]


_OPTIONS = bergamot.ResponseOptions(alignment=False, qualityScores=False, HTML=False)


def _translate(texts: list[str], source: str, target: str) -> list[str]:
    todo = [i for i, t in enumerate(texts) if t.strip()]
    out = list(texts)
    if not todo or source == target:
        return out
    batch = [texts[i] for i in todo]
    steps = [(source, target)] if PIVOT in (source, target) else [(source, PIVOT), (PIVOT, target)]
    for s, t in steps:
        model = _model(s, t)
        batch = [
            r.target.text for r in _service.translate(model, bergamot.VectorString(batch), _OPTIONS)
        ]
    for i, text in zip(todo, batch, strict=True):
        out[i] = text
    return out


# --- identifying languages -----------------------------------------------------------------

# fastText's labels where Firefox's codes differ.
_LID_CODES = {"zh": "zh-Hans", "no": "nb", "iw": "he", "fil": "tl"}
# Signs only traditional Chinese uses (a few common ones): tell zh-Hant apart.
_TRADITIONAL = set("們這個來說會對時國過還後點學東車長門開關問電話員")


def _detect(text: str) -> dict[str, Any]:
    sample = " ".join(text.split())[:2000]
    if not sample:
        return {"language": "", "confidence": 0.0}
    labels, scores = _lid.predict(sample, k=1)
    code = labels[0].replace("__label__", "") if labels else ""
    code = _LID_CODES.get(code, code)
    if code == "zh-Hans" and sum(c in _TRADITIONAL for c in sample) >= 2:
        code = "zh-Hant"
    return {"language": code, "confidence": round(float(scores[0]) * 100, 1) if scores else 0.0}


# --- the API ---------------------------------------------------------------------------------


class DetectIn(BaseModel):
    q: str | list[str]


class TranslateIn(BaseModel):
    q: str | list[str]
    source: str
    target: str


@app.get("/health")
def health() -> dict[str, bool]:
    return {"ok": True}


@app.get("/languages")
def languages() -> dict[str, Any]:
    try:
        available = _available()
    except Exception as exc:  # noqa: BLE001 — offline: what is installed still works
        log.info("model catalogue unavailable: %s", exc)
        available = []
    with _lock:
        return {
            "installed": _installed(),
            "installing": {k: dict(v) for k, v in _installing.items()},
            "pending": [c for c in _pending() if c not in _installing],
            "available": available,
            "errors": dict(_errors),
        }


@app.post("/languages/{code}", status_code=202)
def install(code: str) -> dict[str, Any]:
    if not _CODE.fullmatch(code) or code == PIVOT:
        raise HTTPException(400, "not a language code")
    try:
        _release_files(code, PIVOT)
        _release_files(PIVOT, code)
    except LookupError as exc:
        raise HTTPException(404, str(exc)) from exc
    except Exception as exc:  # noqa: BLE001 — Mozilla out of reach: queue it anyway
        log.info("model catalogue unavailable, %s queued: %s", code, exc)
    _ask_for(code)
    return {"started": True}


@app.delete("/languages/{code}")
def remove(code: str) -> dict[str, bool]:
    if code == PIVOT:
        raise HTTPException(400, "English is always installed")
    with _lock:
        _write_state([c for c in _installed() if c != code], [c for c in _pending() if c != code])
        _errors.pop(code, None)
        for pair in (f"{code}-{PIVOT}", f"{PIVOT}-{code}"):
            _loaded.pop(pair, None)
            shutil.rmtree(MODELS / pair, ignore_errors=True)
    return {"removed": True}


@app.post("/detect")
def detect(body: DetectIn) -> Any:
    if isinstance(body.q, str):
        return [_detect(body.q)]
    return [_detect(q) for q in body.q]


@app.post("/translate")
def translate(body: TranslateIn) -> dict[str, Any]:
    texts = [body.q] if isinstance(body.q, str) else body.q
    source = body.source
    if source == "auto":
        source = _detect(" ".join(texts))["language"]
    installed = _installed()
    for code in (source, body.target):
        if code not in installed:
            raise HTTPException(400, f"{code or 'this language'} is not installed")
    out = _translate(texts, source, body.target)
    return {"translatedText": out[0] if isinstance(body.q, str) else out, "source": source}