local-translator/local_translator/tomie_full.py

"""Resumable bulk translation and offline scenario patch for Tomie."""
import argparse
from copy import deepcopy
import hashlib
import json
from pathlib import Path
import re
import time

import httpx

from .benchmark import write_json
from .models import Project
from .protection import protect
from .tomie import DATA, bundle_crc, digest, ensure_api, inside, original_sources, patch_catalog, patch_lock, snapshot, verify_bundle
from .unity_scripts import ScriptCodec, script_objects

FIELDS = {"PrintText": "Text", "AppendText": "Text", "AddChoice": "ChoiceSummary",
          "ShowToastUI": "Text", "InputCustomVariable": "Summary"}


def text_slots(tree):
    """Yield only known visible text parameters; command IDs/expressions stay intact."""
    for index, ref in enumerate(tree["references"]["records"]):
        field = FIELDS.get(ref["type"]["class"])
        if field is None:
            continue
        parameter = ref["data"][field]
        if parameter.get("hasValue") and isinstance(parameter.get("value"), str):
            if parameter["value"].strip():
                yield index, field, ("value",), parameter["value"]
        dynamic = parameter.get("dynamicValue", {})
        if dynamic.get("ValueText", "").strip():
            yield index, field, ("dynamicValue", "ValueText"), dynamic["ValueText"]


def split_text(text, limit=1000):
    """Split long prose at sentence/space boundaries, never inside protected syntax."""
    protected = protect(text, Project(name="split"))
    forbidden = []
    # Use original token lengths to recover protected spans independently of nonce.
    cursor = 0
    for token in protected.tokens.values():
        start = text.index(token, cursor)
        forbidden.append((start, start + len(token)))
        cursor = start + len(token)
    pieces, start = [], 0
    while len(text) - start > limit:
        end = start + limit
        boundaries = [m.end() for m in re.finditer(r"[.!?][ \t]+|\n|[ \t]+", text[start:end])]
        candidates = [start + point for point in boundaries
                      if not any(a < start + point < b for a, b in forbidden)]
        if not candidates:
            # A long protected expression must remain one piece.
            end = next((b for a, b in forbidden if a <= end < b), end)
        else:
            end = candidates[-1]
        pieces.append(text[start:end])
        start = end
    pieces.append(text[start:])
    assert "".join(pieces) == text
    return pieces


def extract(config, root, output):
    from .unity_stream import load
    game = (root / config["game"]).resolve()
    source = game / DATA / "data.unity3d"
    sources, _ = original_sources(game, [source])
    source_hash = digest(sources[source])
    if (output / "dataset.json").exists():
        saved = json.loads((output / "dataset.json").read_text(encoding="utf-8"))
        if saved["source_sha256"] != source_hash:
            raise ValueError("Game source changed; use a different full_output directory")
        return saved
    env = load(sources[source])
    codec = ScriptCodec(game)
    records, scripts = [], []
    with codec.active():
        for obj in script_objects(env):
            tree = codec.read(obj)
            scripts.append({"asset": obj.assets_file.name, "path_id": obj.path_id, "name": tree["m_Name"],
                            "sha256": hashlib.sha256(obj.get_raw_data()).hexdigest()})
            for ref_index, field, slot, text in text_slots(tree):
                ref = tree["references"]["records"][ref_index]
                line = ref["data"]["playbackSpot"]["lineIndex"]
                records.append({"asset": obj.assets_file.name, "path_id": obj.path_id,
                                "scene": tree["m_Name"], "line": line, "ref": ref_index,
                                "field": field, "slot": list(slot), "text": text,
                                "kind": ref["type"]["class"]})
            print(f"Проверен сценарий {tree['m_Name']}: {len(records)} строк", flush=True)
    # Preserve scenario and source line order for context in model requests.
    records.sort(key=lambda r: (r["scene"] != "Opening", r["scene"], r["line"], r["ref"]))
    items, ids = [], {}
    for record in records:
        record["parts"] = []
        for piece in split_text(record["text"]):
            if piece not in ids:
                identity = f"t{len(items):06d}"
                ids[piece] = identity
                items.append({"id": identity, "text": piece})
            record["parts"].append(ids[piece])
    dataset = {"version": 1, "source_sha256": source_hash, "scripts": scripts,
               "records": records, "items": items, "scope": "scenario dialogue, choices, toast and input summaries"}
    write_json(output / "dataset.json", dataset)
    print(f"Извлечено {len(records)} строк, {len(items)} уникальных фрагментов, "
          f"{sum(len(i['text']) for i in items):,} символов.", flush=True)
    return dataset


def chunks(items, count):
    return [items[start:start + count] for start in range(0, len(items), count)]


def request(client, method, path, **kwargs):
    result = client.request(method, path, **kwargs)
    result.raise_for_status()
    return result.json()


def translate(config, root, output, dataset, retry=False):
    base = config["api_base"].rstrip("/")
    ensure_api(root, base)
    project_config = json.loads((root / config["project_file"]).read_text(encoding="utf-8-sig"))
    project_config["name"] = "Tomie — full scenario translation"
    part_size = config.get("full_part_size", 400)
    if not 1 <= part_size <= 500:
        raise ValueError("full_part_size must be 1..500")
    batch_size = config.get("full_batch_size", 20)
    max_chars = config.get("full_max_batch_chars", 2000)
    plan = {"dataset": digest(output / "dataset.json"), "base": base, "project": project_config,
            "part_size": part_size, "batch_size": batch_size, "max_batch_chars": max_chars}
    state_path = output / "state.json"
    with httpx.Client(base_url=base, timeout=45, trust_env=False) as client:
        if state_path.exists():
            state = json.loads(state_path.read_text(encoding="utf-8"))
            if state["plan"] != plan:
                raise ValueError("Saved run uses different input/settings; use a new full_output directory")
        else:
            project = request(client, "POST", "/projects", json=project_config)
            state = {"plan": plan, "project_id": project["id"], "parts": [], "started": time.time()}
            write_json(state_path, state)
        parts = chunks(dataset["items"], part_size)
        completed, errors = 0, 0
        for index, items in enumerate(parts):
            result_path = output / "parts" / f"{index + 1:04d}.json"
            if index < len(state["parts"]):
                job = request(client, "GET", "/jobs/" + state["parts"][index])
                if retry and job["status"] in {"completed_with_errors", "failed"}:
                    # A previously oversized model response must not be retried
                    # with the same batch. Completed rows come from project cache.
                    old_id = job["id"]
                    job = request(client, "POST", f"/projects/{state['project_id']}/jobs", json={
                        "items": items, "batch_size": 1, "max_batch_chars": max_chars, "retries": 2})
                    state.setdefault("replaced_jobs", []).append(old_id)
                    state.setdefault("single_item_parts", []).append(index)
                    state["parts"][index] = job["id"]
                    write_json(state_path, state)
            else:
                job = request(client, "POST", f"/projects/{state['project_id']}/jobs", json={
                    "items": items, "batch_size": batch_size, "max_batch_chars": max_chars, "retries": 2})
                state["parts"].append(job["id"])
                write_json(state_path, state)
            previous = None
            while True:
                progress = job["progress"]
                done = progress.get("completed", 0)
                failed = progress.get("failed", 0)
                report = {"status": "running", "part": index + 1, "parts": len(parts),
                          "completed": completed + done, "failed": errors + failed,
                          "total": len(dataset["items"]), "job_id": job["id"],
                          "elapsed_seconds": round(time.time() - state["started"], 1)}
                write_json(output / "progress.json", report)
                marker = (job["status"], done, failed)
                if marker != previous:
                    write_json(result_path, job)
                    print(f"Часть {index + 1}/{len(parts)}: {done}/{len(items)}, ошибок {failed}; "
                          f"всего {completed + done}/{len(dataset['items'])}", flush=True)
                    previous = marker
                if job["status"] not in {"queued", "running"}:
                    if job["status"] in {"failed", "completed_with_errors"} and index not in state.get("single_item_parts", []):
                        old_id = job["id"]
                        job = request(client, "POST", f"/projects/{state['project_id']}/jobs", json={
                            "items": items, "batch_size": 1, "max_batch_chars": max_chars, "retries": 2})
                        state.setdefault("replaced_jobs", []).append(old_id)
                        state.setdefault("single_item_parts", []).append(index)
                        state["parts"][index] = job["id"]
                        write_json(state_path, state)
                        print(f"Часть {index + 1}: повтор незавершённых строк по одной, готовые берутся из кэша", flush=True)
                        continue
                    retries = state.setdefault("automatic_retries", {})
                    if job["status"] in {"failed", "completed_with_errors"} and retries.get(str(index), 0) < 2:
                        retries[str(index)] = retries.get(str(index), 0) + 1
                        write_json(state_path, state)
                        print(f"Повторяю только незавершённые строки части {index + 1}: {job.get('error') or failed}", flush=True)
                        time.sleep(3)
                        job = request(client, "POST", "/jobs/" + job["id"] + "/retry", json={})
                        continue
                    break
                time.sleep(3)
                job = request(client, "GET", "/jobs/" + job["id"])
            write_json(result_path, job)
            completed += done
            errors += failed + (len(items) - done - failed)
            if job["status"] == "failed" and not done:
                # Provider outage must not burn retry budgets for the entire game.
                raise RuntimeError(f"Provider failed in part {index + 1}; restart with --retry after fixing it")
        status = "completed" if errors == 0 else "completed_with_errors"
        write_json(output / "progress.json", {**report, "status": status, "completed": completed, "failed": errors})
        if errors:
            raise RuntimeError(f"Не переведено {errors} фрагментов. Повторный запуск с --retry повторит только ошибки.")
        return load_translations(output, dataset)


def load_translations(output, dataset):
    translated = {}
    source = {item["id"]: item["text"] for item in dataset["items"]}
    for path in sorted((output / "parts").glob("*.json")):
        job = json.loads(path.read_text(encoding="utf-8"))
        for item in job["items"]:
            if item["status"] == "completed":
                if source.get(item["id"]) != item["source"]:
                    raise ValueError("Saved translation does not match source dataset")
                translated[item["id"]] = item["translation"]
    if set(translated) != set(source):
        raise ValueError("Translation is incomplete; refusing to build a partial patch")
    return translated


def build(config, root, output, dataset, translated, ui=None, extra=None, extra_translated=None):
    """Add translated scenarios to the verified, previously prepared UI package."""
    from .unity_stream import load, resource_hashes, save, streaming_unity
    game = (root / config["game"]).resolve()
    latest = root / config["output"] / "latest.json"
    ui = ui or Path(json.loads(latest.read_text(encoding="utf-8"))["package"])
    ui_manifest = json.loads((ui / "manifest.json").read_text(encoding="utf-8"))
    if ui_manifest.get("scenario_translation"):
        raise ValueError("Latest package already contains scenarios; select a fresh UI package")
    for record in ui_manifest["files"]:
        if digest(inside(ui / "payload", record["path"])) != record["after"]:
            raise ValueError("UI package checksum mismatch")
    source = game / DATA / "data.unity3d"
    sources, _ = original_sources(game, [source])
    if digest(sources[source]) != dataset["source_sha256"]:
        raise ValueError("Game source does not match extracted scenarios")
    env = load(ui / "payload" / DATA / "data.unity3d")
    before = snapshot(env)
    resources_before = resource_hashes(next(iter(env.files.values())))
    codec = ScriptCodec(game)
    records = {}
    for record in dataset["records"]:
        records.setdefault((record["asset"], record["path_id"]), []).append(record)
    project = Project.model_validate(json.loads((root / config["project_file"]).read_text(encoding="utf-8-sig")))
    allowed, changed, expected = set(), 0, {}
    next_item = {item["id"]: item["text"] for item in dataset["items"]}
    with codec.active():
        for obj in script_objects(env):
            key = (obj.assets_file.name, obj.path_id)
            tree = codec.read(obj)
            for record in records.pop(key, []):
                param = tree["references"]["records"][record["ref"]]["data"][record["field"]]
                for field in record["slot"][:-1]:
                    param = param[field]
                field = record["slot"][-1]
                if param[field] != record["text"]:
                    raise ValueError("Scenario text changed since extraction")
                # Restore exact whitespace at split boundaries, which LLMs may trim.
                pieces = []
                for identity in record["parts"]:
                    source_piece = next_item[identity]
                    target = translated[identity]
                    pieces.append(re.match(r"\s*", source_piece)[0] + target.strip() + re.search(r"\s*$", source_piece)[0])
                target = "".join(pieces)
                if list(protect(record["text"], project).tokens.values()) != list(protect(target, project).tokens.values()):
                    raise ValueError("Protected syntax changed in translated scenario")
                param[field] = target
                changed += 1
            obj.save_typetree(tree, nodes=codec.node("Elringus.Naninovel.Runtime", "Naninovel.Script"))
            expected[key] = hashlib.sha256(json.dumps(tree, sort_keys=True, ensure_ascii=False).encode("utf-8")).hexdigest()
            allowed.add(key)
    if records:
        raise ValueError("Missing scenario objects")
    package = output / "packages" / ("full-" + str(time.time_ns()))
    package.mkdir(parents=True, exist_ok=True)
    import shutil
    extra_files, extra_allowed = [], {"main": set(), "bundle": set()}
    if extra:
        from .tomie_extra import apply_extra
        bundle = load(ui / "payload" / extra["bundle"])
        bundle_before = snapshot(bundle)
        bundle_resources_before = resource_hashes(next(iter(bundle.files.values())))
        extra_allowed, extra_files = apply_extra(env, bundle, extra, extra_translated, game, package / "payload")
        allowed |= extra_allowed["main"]
    for record in ui_manifest["files"]:
        target = inside(package / "payload", record["path"])
        target.parent.mkdir(parents=True, exist_ok=True)
        if record["path"] == (DATA / "data.unity3d").as_posix():
            save(next(iter(env.files.values())), target)
            with streaming_unity():
                verification = verify_bundle(target, before, allowed)
        elif extra and record["path"] == extra["bundle"]:
            new_bundle_crc = save(next(iter(bundle.files.values())), target)
            with streaming_unity():
                verify_bundle(target, bundle_before, extra_allowed["bundle"])
            if resource_hashes(next(iter(load(target).files.values()))) != bundle_resources_before:
                raise ValueError("Addressable media resource bytes changed")
        else:
            shutil.copyfile(inside(ui / "payload", record["path"]), target)
    if extra:
        bundle_payload = inside(package / "payload", extra["bundle"])
        catalog_path = next(r["path"] for r in ui_manifest["files"] if r["path"].endswith("/catalog.json"))
        document = json.loads(inside(ui / "payload", catalog_path).read_text(encoding="utf-8"))
        old_hash = ui_manifest["addressables_after"]["m_Hash"]
        with bundle_payload.open("rb") as stream:
            new_hash = hashlib.file_digest(stream, "md5").hexdigest()
        document, _, new_options = patch_catalog(document, old_hash, new_bundle_crc, bundle_payload.stat().st_size, new_hash)
        write_json(inside(package / "payload", catalog_path), document)
    # Read every packed scenario again, checking the entire object, including all
    # commands, references, hashes, conditions and line numbers against intent.
    packed = load(package / "payload" / DATA / "data.unity3d")
    if resource_hashes(next(iter(packed.files.values()))) != resources_before:
        raise ValueError("Game media resource bytes changed")
    with codec.active():
        for obj in script_objects(packed):
            key = (obj.assets_file.name, obj.path_id)
            tree = codec.read(obj)
            fingerprint = hashlib.sha256(json.dumps(tree, sort_keys=True, ensure_ascii=False).encode("utf-8")).hexdigest()
            if expected.pop(key, None) != fingerprint:
                raise ValueError("Packed scenario differs from intended text-only changes")
    if expected:
        raise ValueError("Packed scenarios are missing")
    manifest = deepcopy(ui_manifest)
    manifest["id"] = package.name
    for record in manifest["files"]:
        record["after"] = digest(inside(package / "payload", record["path"]))
    manifest["scenario_translation"] = {"scripts": len(dataset["scripts"]), "strings": changed, "verification": verification}
    manifest["files"].extend(extra_files)
    if extra:
        manifest["extra_translation"] = {"strings": len(extra["records"]), "stream_files": len(extra_files)}
        manifest["addressables_after"] = new_options
    manifest["scope"] = "UI, managed text, compiled scenarios and dynamic dialogue templates; no image text or in-game verification"
    write_json(package / "manifest.json", manifest)
    write_json(latest, {"package": str(package)})
    write_json(output / "built.json", {"package": str(package)})
    print(f"Готов полный патч сценариев: {package}", flush=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=Path("tomie.json"))
    parser.add_argument("action", choices=("extract", "translate", "build", "run", "status"))
    parser.add_argument("--retry", action="store_true")
    args = parser.parse_args()
    root = args.config.resolve().parent
    config = json.loads(args.config.read_text(encoding="utf-8-sig"))
    output = root / config.get("full_output", "data/tomie/full-translation")
    output.mkdir(parents=True, exist_ok=True)
    if args.action == "status":
        if not (output / "progress.json").exists():
            print("Ещё не запущен")
            return
        progress = json.loads((output / "progress.json").read_text(encoding="utf-8"))
        print(f"{progress['status']}: часть {progress['part']}/{progress['parts']}; "
              f"переведено {progress['completed']}/{progress['total']}, ошибок {progress['failed']}.")
        print(f"Результаты: {output}\nПовторный запуск translate_all_tomie.bat продолжит сохранённую работу.")
        return
    with patch_lock(output):
        dataset = extract(config, root, output)
        ui_state = output / "ui-package.json"
        if not ui_state.exists():
            write_json(ui_state, json.loads((root / config["output"] / "latest.json").read_text(encoding="utf-8")))
        ui = Path(json.loads(ui_state.read_text(encoding="utf-8"))["package"])
        from .tomie_extra import extract_extra
        extra_output = output / "extra"
        if args.action == "extract":
            return
        extra = extract_extra(config, root, extra_output, ui)
        if args.action in {"translate", "run"}:
            translate(config, root, output, dataset, args.retry)
            translate(config, root, extra_output, extra, args.retry)
        if args.action in {"build", "run"}:
            translated = load_translations(output, dataset)
            extra_translated = load_translations(extra_output, extra)
            build(config, root, output, dataset, translated, ui, extra, extra_translated)


if __name__ == "__main__":
    main()