"""MarianMT translate worker — runs in separate process.

Usage: python -m src.math_solver.translate_worker
Reads text from stdin, translates Chinese → English, outputs to stdout.
"""
import json
import re
import sys
import os

os.environ["PYTHONIOENCODING"] = "utf-8"
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
if sys.stdout.encoding != "utf-8":
    sys.stdout.reconfigure(encoding="utf-8")
if sys.stdin.encoding != "utf-8":
    sys.stdin.reconfigure(encoding="utf-8")


def main():
    text = sys.stdin.read().strip()
    if not text:
        print(json.dumps({"text": ""}))
        return

    from transformers import MarianMTModel, MarianTokenizer

    model_name = "Helsinki-NLP/opus-mt-zh-en"
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)

    lines = text.split("\n")
    translated_lines = []

    for line in lines:
        line = line.strip()
        if not line:
            continue
        has_chinese = bool(re.search(r'[\u4e00-\u9fff]', line))
        if has_chinese:
            inputs = tokenizer(line, return_tensors="pt", padding=True, truncation=True, max_length=512)
            outputs = model.generate(**inputs)
            en_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
            translated_lines.append(en_text)

    print(json.dumps({"text": "\n".join(translated_lines)}))


if __name__ == "__main__":
    main()
