#!/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 — the full URL
including the access key (printed at server startup; the admin also finds
the key in the site settings, GET /_api/settings -> ``translate_keys``) —
and announces the languages the
model CAN translate (capabilities). The server dispatches one single-item
job at a time per connection, offered only in its configured target
languages (``Data.translate_langs``) ∩ the announced capabilities; a
dropped connection's in-flight item is simply re-offered
(docs/localization.md). For parallelism, run multiple instances.

Seed-X-PPO-7B (bf16, ~15 GB) is the only supported model.

Usage:
    uv run scripts/translator.py ws://localhost:8410/_translate/KEY
    uv run scripts/translator.py wss://example.com/_translate/KEY
"""

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 masked Markdown (pagerite/masking.py: ⟦N⟧ sentinels
#: stand in for code, URLs, placeholders...); Seed-X has no system prompt,
#: so the instruction goes in-line.
NOTE = ", preserving all Markdown formatting and keeping every ⟦N⟧ token exactly unchanged"


# The wire structs below 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"):
    """Client greeting on connect: the language codes its model CAN produce
    (capabilities). The server offers jobs only in the intersection with
    its wanted target languages."""

    langs: list[str]


class Job(msgspec.Struct, tag="job"):
    """Server push: ONE fragment to translate. Exactly one job is in flight
    per connection — the next arrives only after this one's Result."""

    lang: str
    key: bytes  #: 9-byte chunk hash (base64 in the JSON frame)
    text: str  #: masked (pagerite/masking.py): the ⟦N⟧ tokens must survive verbatim
    path: str  #: article it came from ("" = front page), no leading slash
    kind: str  #: "chunk" | "title"


class Result(msgspec.Struct, tag="result"):
    """Client reply: the translation of the connection's current Job
    (must match its lang and key exactly)."""

    lang: str
    key: bytes
    text: str


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 the job's one fragment and send the result back."""
    lang_name = SEED_X_NAMES[job.lang].capitalize()
    # Deliberately blocking: nothing else needs the loop while the job is
    # being answered, and the reconnect loop recovers a dropped connection
    # (the in-flight item is simply re-offered).
    text, tokens, dt = seed_x_chunk(tokenizer, model, job.text, lang_name, job.lang, note=NOTE)
    print(f"[{job.lang} {job.kind} {job.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, key=job.key, text=text)).decode())


async def serve(url: str, tokenizer, model) -> None:
    """Connect, announce capabilities, 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=sorted(SEED_X_NAMES))).decode())
                print(f"[connected; announced {len(SEED_X_NAMES)} language capabilities]",
                      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 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():
    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")
    args = p.parse_args()
    if not args.url.startswith(("ws://", "wss://")):
        p.error("url must start with ws:// or wss://")

    tokenizer, model = load_seed_x()  # once, before the (re)connect loop
    asyncio.run(serve(args.url, tokenizer, model))


if __name__ == "__main__":
    main()