#!/usr/bin/env python3
"""Translate the hard-wrapped and unwrapped sample into Simplified Chinese with a local model.

Conditions:
  wrapped     source-wrapped.txt as it is (lines broken at 66 characters)
  unwrapped   source-unwrapped.txt (one paragraph per line, same words)
  chunked     source-wrapped.txt cut every 12 lines, as a line-count splitter would,
              each chunk sent as a separate request and the replies joined with a newline

All three use the same prompt. Usage:
    python3 run_llm_txt_test.py --runs 3
Outputs: outputs/<condition>-<run>.txt

Written by BookTranslator (https://www.booktranslator.app). Free to reuse.
"""

import argparse
import json
import urllib.request
from pathlib import Path

HERE = Path(__file__).parent
URL = "http://localhost:1234/v1/chat/completions"
MODEL = "qwen/qwen3.6-27b"
PROMPT = (
    "Translate the following plain-text file into Simplified Chinese. "
    "Return only the translated text, with no commentary.\n\n"
)


def ask(text):
    body = {
        "model": MODEL,
        "messages": [{"role": "user", "content": PROMPT + text}],
        "temperature": 0.3,
        "reasoning_effort": "none",
    }
    req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=900) as r:
        return json.load(r)["choices"][0]["message"]["content"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--conditions", default="wrapped,unwrapped,chunked")
    args = ap.parse_args()
    out = HERE / "outputs"
    out.mkdir(exist_ok=True)
    for condition in args.conditions.split(","):
        name = "wrapped" if condition == "chunked" else condition
        source = (HERE / f"source-{name}.txt").read_text(encoding="utf-8")
        for run in range(1, args.runs + 1):
            if condition == "chunked":
                lines = source.split("\n")
                chunks = ["\n".join(lines[i:i + 12]) for i in range(0, len(lines), 12)]
                reply = "\n".join(ask(c).strip("\n") for c in chunks if c.strip()) + "\n"
            else:
                reply = ask(source).strip("\n") + "\n"
            (out / f"{condition}-{run}.txt").write_text(reply, encoding="utf-8")
            print(f"{condition}-{run}: {len(reply)} chars")


if __name__ == "__main__":
    main()