"""Reproducible local-LLM long-document translation test (LM Studio OpenAI-compatible API).

Source: Kafka, Die Verwandlung, Project Gutenberg #22367 (public domain).
Download it next to this script first:
  curl -L https://www.gutenberg.org/cache/epub/22367/pg22367.txt -o verwandlung.txt
Then start the LM Studio server with the model loaded and run: python3 local_llm_translation_test.py A B C
Thinking disabled with reasoning_effort="none" ("/no_think" and chat_template_kwargs did not stop reasoning in LM Studio).
Conditions:
  A  chunked, bare prompt, no glossary, no context
  B  chunked, glossary + previous source paragraph as read-only context
  C  one request with the whole excerpt
Outputs raw JSON per request plus results.json with timings and checks.
"""
import json, re, sys, time, urllib.request
from pathlib import Path

HERE = Path(__file__).parent
API = "http://localhost:1234/v1/chat/completions"
MODEL = "qwen/qwen3.6-27b"
TARGET_WORDS = 700
EXCERPT_WORDS = 3000

raw = (HERE / "verwandlung.txt").read_text(encoding="utf-8")
body = raw.split("*** START OF THE PROJECT GUTENBERG EBOOK DIE VERWANDLUNG ***")[1]
body = body[body.index("Als Gregor Samsa"):]
paras = [re.sub(r"\s+", " ", p).strip() for p in re.split(r"\n\s*\n", body)]
paras = [p for p in paras if p]

excerpt, n = [], 0
for p in paras:
    if n >= EXCERPT_WORDS:
        break
    excerpt.append(p)
    n += len(p.split())

chunks, cur, cw = [], [], 0
for p in excerpt:
    if cur and cw + len(p.split()) > TARGET_WORDS:
        chunks.append(cur)
        cur, cw = [], 0
    cur.append(p)
    cw += len(p.split())
if cur:
    chunks.append(cur)

GLOSSARY = """Gregor Samsa = Gregor Samsa
Grete = Grete
Prokurist = chief clerk
Chef = boss
Ungeziefer = vermin
Geschäftsreisender / Reisender = traveling salesman
Kommis = junior clerk"""

BASE = ("Translate the German text into English. Keep exactly one output paragraph per input "
        "paragraph, separated by a blank line. Output only the translation.")


def call(system, user):
    payload = {"model": MODEL, "temperature": 0.2, "max_tokens": 12000, "reasoning_effort": "none",
               "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}]}
    req = urllib.request.Request(API, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
    t = time.time()
    with urllib.request.urlopen(req, timeout=3600) as r:
        out = json.load(r)
    out["_wall_seconds"] = round(time.time() - t, 1)
    return out


def text_of(resp):
    t = resp["choices"][0]["message"]["content"] or ""
    return re.sub(r"<think>.*?</think>", "", t, flags=re.S).strip()


TERMS = {"Prokurist": ["chief clerk", "manager", "head clerk", "procurator", "authorized", "office manager", "senior clerk", "chief manager"],
         "Ungeziefer": ["vermin", "insect", "bug", "beetle", "creature", "pest"]}


def check(src_paras, out_text):
    out_paras = [p for p in re.split(r"\n\s*\n", out_text) if p.strip()]
    low = out_text.lower()
    return {"src_paragraphs": len(src_paras), "out_paragraphs": len(out_paras),
            "src_words": sum(len(p.split()) for p in src_paras), "out_words": len(out_text.split()),
            "renderings": {k: {v: low.count(v) for v in vs if low.count(v)} for k, vs in TERMS.items()
                           if any(k in p for p in src_paras)}}


results = {"model": MODEL, "excerpt_words": n, "excerpt_paragraphs": len(excerpt),
           "chunks": [len(c) for c in chunks], "runs": {}}
conds = sys.argv[1:] or ["A", "B", "C"]
for cond in conds:
    runs = []
    if cond == "C":
        groups = [excerpt]
    else:
        groups = chunks
    for i, g in enumerate(groups):
        system = BASE
        user = "\n\n".join(g)
        if cond == "B":
            system = BASE + "\nUse this glossary exactly:\n" + GLOSSARY
            prev = groups[i - 1][-1] if i else ""
            if prev:
                user = ("Context (already translated, do NOT translate):\n" + prev +
                        "\n\nText to translate:\n" + user)
        resp = call(system, user)
        (HERE / f"raw-{cond}-{i}.json").write_text(json.dumps(resp, ensure_ascii=False, indent=1))
        out = text_of(resp)
        (HERE / f"out-{cond}-{i}.txt").write_text(out)
        u = resp.get("usage", {})
        tps = (resp.get("stats") or {}).get("tokens_per_second")
        runs.append({"chunk": i, "wall_s": resp["_wall_seconds"], "usage": u, "tokens_per_second": tps,
                     "finish_reason": resp["choices"][0].get("finish_reason"), **check(g, out)})
        print(cond, i, runs[-1], flush=True)
    results["runs"][cond] = runs
    prev = json.loads((HERE / "results.json").read_text()) if (HERE / "results.json").exists() else {}
    prev.update({k: v for k, v in results.items() if k != "runs"})
    prev.setdefault("runs", {}).update({cond: runs})
    (HERE / "results.json").write_text(json.dumps(prev, ensure_ascii=False, indent=1))