#!/usr/bin/env python3
# /// script
# requires-python = ">=3.14"
# dependencies = [
#     "accelerate>=1.14.0",
#     "msgspec>=0.19.0",
#     "torch>=2.13.0",
#     "tracerite>=2.6.5",
#     "transformers>=5.16.1",
#     "websockets>=15.0.1",
# ]
# ///
"""Pagerite translator service: translate site content with Seed-X-PPO-7B.

Connects to a Pagerite server's translator WebSocket (``/_translate/<key>`` —
deliberately outside /_api, the key is the access control; find it in the
site settings, GET /_api/settings -> ``translate_key``), announces the
languages it handles and translates whatever the server pushes: everything
pending on connect, then deltas as the content changes. Results go back per
job and the server stores them (docs/localization.md).

Seed-X-PPO-7B (bf16, ~15 GB) is the only supported model; its 28 languages
are the ceiling of what a site can announce through this script.

Usage:
    uv run scripts/translator.py ws://localhost:8080 --key "..." --to finnish
    uv run scripts/translator.py wss://example.com --key "..." --to fi german es
"""

import argparse
import asyncio
import sys
import time

import msgspec
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import tracerite
import websockets

tracerite.load()

SEED_X = "ByteDance-Seed/Seed-X-PPO-7B"

# Seed-X language tags (appended to the prompt; required by its PPO training)
SEED_X_TAGS = {
    "arabic": "ar", "chinese": "zh", "czech": "cs", "danish": "da",
    "dutch": "nl", "english": "en", "finnish": "fi", "french": "fr",
    "german": "de", "greek": "el", "hungarian": "hu", "indonesian": "id",
    "italian": "it", "japanese": "ja", "korean": "ko", "malay": "ms",
    "norwegian": "no", "persian": "fa", "polish": "pl", "portuguese": "pt",
    "romanian": "ro", "russian": "ru", "spanish": "es", "swedish": "sv",
    "thai": "th", "turkish": "tr", "ukrainian": "uk", "vietnamese": "vi",
}
SEED_X_NAMES = {v: k for k, v in SEED_X_TAGS.items()}

#: The fragments are Markdown; Seed-X has no system prompt, so it goes in-line.
NOTE = ", preserving all Markdown formatting, URLs and code exactly unchanged"


# The message structs below duplicate pagerite/translate.py 1:1: 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"):
    """Client greeting on connect: the target languages it handles."""

    langs: list[str]


class TransItem(msgspec.Struct):
    """One fragment to translate: original Markdown (or a node title)."""

    key: bytes  #: 9-byte chunk hash (base64 in the JSON frame)
    text: str
    path: str  #: article it came from ("" = front page), no leading slash
    kind: str  #: "chunk" | "title"


class Job(msgspec.Struct, tag="job"):
    """Server push: pending items for one language."""

    lang: str
    items: list[TransItem]


class TransResult(msgspec.Struct):
    """One translated fragment."""

    key: bytes
    text: str


class Result(msgspec.Struct, tag="result"):
    """Client reply: a batch of translations for one language."""

    lang: str
    items: list[TransResult]


def load_seed_x():
    t0 = time.monotonic()
    tokenizer = AutoTokenizer.from_pretrained(SEED_X)
    model = AutoModelForCausalLM.from_pretrained(SEED_X, dtype=torch.bfloat16, device_map="auto")
    print(f"[seed-x loaded in {time.monotonic() - t0:.0f}s]", file=sys.stderr)
    return tokenizer, model


def seed_x_chunk(tokenizer, model, text: str, target_lang: str, tag: str, source_lang: str = "English",
                 note: str = ""):
    """Translate one segment; returns (translation, output_tokens, generation_seconds)."""
    # No chat template on this model; the trailing language tag is required (trans/ style prompt).
    prompt = f"Translate the following {source_lang} text into {target_lang}{note}:\n{text} <{tag}>"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    t0 = time.monotonic()
    out = model.generate(**inputs, max_new_tokens=max(1024, 2 * inputs.input_ids.shape[1]), do_sample=False)
    dt = time.monotonic() - t0
    n = out.shape[1] - inputs.input_ids.shape[1]
    return tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip(), n, dt


async def do_job(ws, job: Job, tokenizer, model) -> None:
    """Translate every item of one job and send the results back as a batch."""
    lang_name = SEED_X_NAMES[job.lang].capitalize()
    results = []
    for i, item in enumerate(job.items, 1):
        # Deliberately blocking: nothing else needs the loop while a job is
        # being answered, and the reconnect loop recovers a dropped connection
        # (results already stored stay stored; only the still-pending items
        # are re-pushed).
        text, tokens, dt = seed_x_chunk(tokenizer, model, item.text, lang_name, job.lang, note=NOTE)
        results.append(TransResult(key=item.key, text=text))
        print(f"[{job.lang} {i}/{len(job.items)} {item.kind} {item.path or '/'}: "
              f"{tokens} tokens in {dt:.1f}s = {tokens / dt:.1f} tok/s]", file=sys.stderr)
    await ws.send(msgspec.json.encode(Result(lang=job.lang, items=results)).decode())


async def serve(url: str, codes: list[str], tokenizer, model) -> None:
    """Connect, announce languages, answer jobs; reconnect with backoff."""
    backoff = 1
    while True:
        try:
            async with websockets.connect(url) as ws:
                backoff = 1
                await ws.send(msgspec.json.encode(Hello(langs=codes)).decode())
                print(f"[connected; translating: {', '.join(codes)}]", file=sys.stderr)
                async for raw in ws:
                    await do_job(ws, msgspec.json.decode(raw, type=Job), tokenizer, model)
        except websockets.exceptions.InvalidHandshake:
            sys.exit("handshake rejected; check --key and the server URL")
        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():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("server", help="server WebSocket base URL, e.g. ws://localhost:8080")
    p.add_argument("--key", required=True, help="translator API key (site settings: translate_key)")
    p.add_argument("--to", required=True, nargs="+",
                   help="target language(s): Seed-X names or codes, e.g. finnish de es")
    args = p.parse_args()

    codes = []
    for lang in args.to:
        low = lang.lower()
        code = SEED_X_TAGS.get(low) or (low if low in SEED_X_NAMES else None)
        if code is None:
            p.error(f"unknown language {lang!r}; supported: {', '.join(sorted(SEED_X_TAGS))}")
        if code not in codes:
            codes.append(code)

    tokenizer, model = load_seed_x()  # once, before the (re)connect loop
    url = f"{args.server.rstrip('/')}/_translate/{args.key}"
    asyncio.run(serve(url, codes, tokenizer, model))


if __name__ == "__main__":
    main()