#!/usr/bin/env python3
# /// script
# requires-python = ">=3.14"
# dependencies = [
#     "httpx>=0.28.1",
#     "msgspec>=0.19.0",
#     "websockets>=15.0.1",
# ]
# ///
"""Pagerite LLM translator service: translate site content with an instruct
LLM that handles Markdown natively (docs/llm-translation.md).

Same channel as scripts/translator.py (Seed-X) — connect to the server's
translator WebSocket URL including its access key, announce capabilities,
answer one job at a time — but speaks the "markdown" and "article" job
modes: fragments and whole pages cross as Markdown, and the server
validates structure (blocks, fences, URLs, placeholders) before storing.

The LLM is reached via an OpenAI Chat Completions endpoint
(base_url + /v1/chat/completions: llama.cpp, hosted APIs) or ollama's
native /api/chat (api="ollama") — ollama's OpenAI endpoint ignores
think:false, which hybrid models need off. Backend quirks (sampling,
num_predict cap, think) live in the config, not in the protocol.

Usage:
    uv run scripts/llm_translator.py ws://localhost:8410/_translate/KEY
    uv run scripts/llm_translator.py wss://example.com/_translate/KEY --config my.json
"""

import argparse
import asyncio
import json
import sys
import time
from pathlib import Path

import httpx
import msgspec
import websockets

#: Shipped defaults, aimed at a local ollama running the structure-proven
#: qwen3.8:27b (docs/llm-translation.md trial evidence). A --config JSON
#: overrides per key, CLI flags override the config.
DEFAULT_CONFIG = {
    "api": "ollama",  # "ollama" (native /api/chat) | "openai" (/v1/chat/completions)
    "base_url": "http://127.0.0.1:11434",
    "model": "qwen3.8:27b",
    "api_key": "",  # openai api only
    "langs": ["de", "es", "fi", "pt", "zh"],  # announced capabilities
    "modes": ["markdown", "article"],
    "temperature": 0.2,
    "top_p": 0.8,
    "top_k": 20,
    "num_ctx": 32768,
    # Generation cap: runaway thinking/generation on a whole-article job
    # burns hours otherwise. num_predict = clamp(src_tokens * ratio, ...).
    "predict_ratio": 2.5,
    "predict_min": 1024,
    "predict_cap": 16384,
    "think": False,  # ollama api only: hybrid models must not think
    "timeout": 10800,
}

LANG_NAMES = {
    "de": "German",
    "es": "Spanish",
    "fi": "Finnish",
    "fr": "French",
    "it": "Italian",
    "ja": "Japanese",
    "ko": "Korean",
    "nl": "Dutch",
    "pl": "Polish",
    "pt": "Portuguese",
    "ru": "Russian",
    "sv": "Swedish",
    "zh": "Simplified Chinese",
}

RULES = """\
Rules:
- Output ONLY the translation, no commentary, no preamble.
- Preserve the Markdown structure exactly: same blocks separated by blank \
lines, same headings (# levels), lists, code fences, images and links.
- Never translate or alter URLs, image destinations, code, or {...} \
placeholders. Image alt texts and link texts ARE translated.
- Do not merge, split, add, drop or reorder blocks."""


def article_prompt(target: str, doc: str) -> str:
    return f"""Translate the following Markdown document into {target}.

{RULES}

```markdown
{doc}
```"""


def block_prompt(target: str, text: str, prev: str, next_: str) -> str:
    prompt = f"""Translate one block of a Markdown document into {target}.

{RULES}
- Translate ONLY the block marked TRANSLATE. The CONTEXT blocks are the \
surrounding document, already translated — terminology and tone \
reference only; never translate or repeat them.
"""
    if prev:
        prompt += f"\nCONTEXT BEFORE (do not translate):\n```markdown\n{prev}\n```\n"
    if next_:
        prompt += f"\nCONTEXT AFTER (do not translate):\n```markdown\n{next_}\n```\n"
    return prompt + f"\nTRANSLATE:\n```markdown\n{text}\n```"


def title_prompt(target: str, title: str, context: str) -> str:
    prompt = f"""Translate the following title into {target}.
Output ONLY the translated title: a single line of plain text, no \
Markdown, no quotes, no commentary, no terminal punctuation unless the \
original has it.
"""
    if context:
        prompt += f"\nThe article it heads begins as follows (context only, do not translate):\n{context}\n"
    return prompt + f"\nTITLE:\n{title}"


# The wire structs duplicate pagerite/translate.py: this script runs in its
# own uv environment and cannot import the server package. The "type" tag
# selects the frame; bytes fields ride as base64.
class Hello(msgspec.Struct, tag="hello"):
    langs: list[str]  #: language codes the model can produce
    model: str = ""
    modes: list[str] = msgspec.field(default_factory=lambda: ["segments"])


class Job(msgspec.Struct, tag="job"):
    """Server push: ONE fragment to translate (next arrives only after the
    Result). markdown/article modes carry a single text — the fragment's /
    the whole page's Markdown."""

    lang: str
    key: bytes
    texts: list[str]
    path: str
    kind: str  #: "chunk" | "title" | "article"
    mode: str = "segments"
    #: markdown mode: [previous, next] block of the served hybrid (target
    #: language); titles: the article's opening. Reference only.
    contexts: list[str] = msgspec.field(default_factory=list)


class Result(msgspec.Struct, tag="result"):
    lang: str
    key: bytes
    texts: list[str]


def unwrap_fence(source: str, out: str) -> str:
    """Strip a whole-output markdown fence the model added around its
    answer (but never when the source itself is fenced)."""
    out = out.strip()
    if (
        not source.lstrip().startswith("```")
        and out.startswith("```")
        and out.endswith("```")
        and len(lines := out.split("\n")) > 2
    ):
        out = "\n".join(lines[1:-1]).strip()
    return out


async def generate(cfg: dict, http: httpx.AsyncClient, prompt: str, src_chars: int) -> tuple[str, int, float]:
    """One chat completion; returns (content, output tokens, seconds)."""
    est = int(src_chars / 3)  # generous token estimate of the source text
    predict = int(
        min(cfg["predict_cap"], max(cfg["predict_min"], est * cfg["predict_ratio"]))
    )
    t0 = time.monotonic()
    if cfg["api"] == "ollama":
        r = await http.post(
            f"{cfg['base_url']}/api/chat",
            json={
                "model": cfg["model"],
                "messages": [{"role": "user", "content": prompt}],
                "stream": False,
                "think": cfg["think"],
                "options": {
                    "temperature": cfg["temperature"],
                    "top_p": cfg["top_p"],
                    "top_k": cfg["top_k"],
                    "num_ctx": cfg["num_ctx"],
                    "num_predict": predict,
                },
            },
        )
        r.raise_for_status()
        d = r.json()
        return d["message"]["content"], d.get("eval_count", 0), time.monotonic() - t0
    headers = {"Authorization": f"Bearer {cfg['api_key']}"} if cfg["api_key"] else {}
    r = await http.post(
        f"{cfg['base_url']}/v1/chat/completions",
        headers=headers,
        json={
            "model": cfg["model"],
            "messages": [{"role": "user", "content": prompt}],
            "temperature": cfg["temperature"],
            "top_p": cfg["top_p"],
            "max_tokens": predict,
        },
    )
    r.raise_for_status()
    d = r.json()
    content = d["choices"][0]["message"]["content"] or ""
    return content, d.get("usage", {}).get("completion_tokens", 0), time.monotonic() - t0


async def do_job(cfg: dict, http: httpx.AsyncClient, ws, job: Job) -> None:
    """Answer one job: build the prompt for its mode, generate, clean up,
    send the Result."""
    target = LANG_NAMES.get(job.lang, job.lang)
    src = job.texts[0]
    if job.mode == "article":
        prompt = article_prompt(target, src)
    elif job.kind == "title":
        prompt = title_prompt(target, src, job.contexts[0] if job.contexts else "")
    else:  # markdown chunk
        prev, next_ = (job.contexts + ["", ""])[:2]
        prompt = block_prompt(target, src, prev, next_)
    out, tokens, dt = await generate(cfg, http, prompt, len(src))
    out = unwrap_fence(src, out)
    if job.kind == "title":
        out = out.split("\n", 1)[0].strip()
    print(
        f"[{job.lang} {job.mode}:{job.kind} {job.path or '/'}: {len(src)} -> "
        f"{len(out)} chars, {tokens} tokens in {dt:.1f}s]",
        file=sys.stderr,
    )
    await ws.send(
        msgspec.json.encode(Result(lang=job.lang, key=job.key, texts=[out])).decode()
    )


async def serve(cfg: dict) -> None:
    """Connect, announce capabilities, answer jobs; reconnect with backoff."""
    url, backoff = cfg["url"], 1
    limits = httpx.Timeout(cfg["timeout"])
    async with httpx.AsyncClient(timeout=limits) as http:
        while True:
            try:
                async with websockets.connect(url) as ws:
                    backoff = 1
                    await ws.send(
                        msgspec.json.encode(
                            Hello(langs=cfg["langs"], model=cfg["model"], modes=cfg["modes"])
                        ).decode()
                    )
                    print(
                        f"[connected; model={cfg['model']}, modes={cfg['modes']}, langs={cfg['langs']}]",
                        file=sys.stderr,
                    )
                    async for raw in ws:
                        await do_job(cfg, http, ws, msgspec.json.decode(raw, type=Job))
            except websockets.exceptions.InvalidHandshake:
                sys.exit("handshake rejected; check the URL (including the key)")
            except (OSError, websockets.exceptions.ConnectionClosed) as e:
                print(
                    f"[connection lost ({e}); reconnecting in {backoff}s]",
                    file=sys.stderr,
                )
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, 60)


def main() -> None:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument(
        "url",
        help="full translator WebSocket URL including the key, "
        "e.g. ws://localhost:8410/_translate/KEY",
    )
    p.add_argument("--config", help="JSON config file (overrides the shipped defaults)")
    p.add_argument("--base-url", help="LLM server root (no path)")
    p.add_argument("--model", help="model string to serve")
    p.add_argument("--api-key", help="API key for openai-api backends")
    p.add_argument("--langs", help="comma-separated announced languages")
    p.add_argument("--modes", help="comma-separated accepted job modes")
    args = p.parse_args()
    if not args.url.startswith(("ws://", "wss://")):
        p.error("url must start with ws:// or wss://")

    cfg = dict(DEFAULT_CONFIG)
    if args.config:
        cfg.update(json.loads(Path(args.config).read_text()))
    for key in ("base_url", "model", "api_key"):
        if getattr(args, key):
            cfg[key] = getattr(args, key)
    if args.langs:
        cfg["langs"] = args.langs.split(",")
    if args.modes:
        cfg["modes"] = args.modes.split(",")
    cfg["url"] = args.url

    try:
        asyncio.run(serve(cfg))
    except (KeyboardInterrupt, asyncio.CancelledError):
        pass


if __name__ == "__main__":
    main()