#!/usr/bin/env python3
"""JSON-lines bridge for Argos Translate and LibreTranslate."""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Callable


TOKEN_RE = re.compile(r"%(?:\d+\$?)?[A-Za-z]|\\[nrt]|<[^>]+>|&(?:#\d+|#x[0-9A-Fa-f]+|\w+);")


def protect(text: str) -> tuple[str, dict[str, str]]:
    tokens: dict[str, str] = {}

    def replace(match: re.Match[str]) -> str:
        key = f"ZXQPH{len(tokens)}QXZ"
        tokens[key] = match.group(0)
        return key

    return TOKEN_RE.sub(replace, text), tokens


def restore(text: str, tokens: dict[str, str]) -> str:
    for key, value in tokens.items():
        text = text.replace(key, value)
    missing = [key for key in tokens if key in text]
    if missing:
        raise RuntimeError(f"translator damaged protected tokens {missing}")
    return text


@dataclass(frozen=True)
class Request:
    texts: list[str]
    source: str
    target: str
    backend: str
    url: str | None
    api_key: str | None


def libre_translator(request: Request) -> Callable[[str], str]:
    if not request.url:
        raise RuntimeError("LibreTranslate requires a URL")
    endpoint = request.url.rstrip("/") + "/translate"

    def translate(text: str) -> str:
        payload: dict[str, object] = {
            "q": text,
            "source": request.source,
            "target": request.target,
            "format": "text",
        }
        if request.api_key:
            payload["api_key"] = request.api_key
        http_request = urllib.request.Request(
            endpoint,
            data=json.dumps(payload).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(http_request, timeout=60) as response:
                body = json.load(response)
        except urllib.error.URLError as error:
            raise RuntimeError(f"LibreTranslate request failed at {endpoint}: {error}") from error
        translated = body.get("translatedText")
        if not isinstance(translated, str):
            raise RuntimeError("LibreTranslate returned no translatedText")
        return translated

    return translate


def argos_translator(request: Request) -> Callable[[str], str]:
    os.environ.setdefault("ARGOS_CHUNK_TYPE", "ARGOSTRANSLATE")
    os.environ.setdefault("ARGOS_DEVICE_TYPE", "cpu")
    try:
        import argostranslate.translate  # type: ignore[import-not-found]
    except ImportError as error:
        raise RuntimeError("Argos Translate is unavailable. Install requirements-translator.txt") from error

    installed = argostranslate.translate.get_installed_languages()
    source = next((language for language in installed if language.code == request.source), None)
    target = next((language for language in installed if language.code == request.target), None)
    if source is None or target is None:
        raise RuntimeError(f"Argos model {request.source} to {request.target} is not installed")
    translation = source.get_translation(target)
    # Android resources are individual UI strings. Avoid network-backed sentence
    # model downloads and pass each protected resource to the installed model.
    class ResourceStringSentencizer:
        @staticmethod
        def split_sentences(text: str) -> list[str]:
            return [text]

    underlying = getattr(translation, "underlying", translation)
    if hasattr(underlying, "sentencizer"):
        underlying.sentencizer = ResourceStringSentencizer()
    return translation.translate


def process(raw: dict[str, object]) -> dict[str, object]:
    request = Request(
        texts=[str(value) for value in raw.get("texts", [])],
        source=str(raw.get("source", "zh")),
        target=str(raw.get("target", "en")),
        backend=str(raw.get("backend", "argos")),
        url=str(raw["url"]) if raw.get("url") else None,
        api_key=str(raw["apiKey"]) if raw.get("apiKey") else None,
    )
    translator = libre_translator(request) if request.backend == "libre" else argos_translator(request)
    output: list[str] = []
    for text in request.texts:
        protected, tokens = protect(text)
        output.append(restore(translator(protected), tokens))
    return {"translations": output}


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--request", help="single JSON request, otherwise stdin is used")
    args = parser.parse_args()
    try:
        raw = json.loads(args.request) if args.request else json.load(sys.stdin)
        json.dump(process(raw), sys.stdout, ensure_ascii=False)
        sys.stdout.write("\n")
        return 0
    except Exception as error:
        print(str(error), file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())