#!/usr/bin/env python3
"""hymt2_batch.py — batch Hy-MT2 translation.

Reads a JSON input listing source strings, runs each through ComfyUI, writes
one JSON output per item under --out-dir. Already-existing outputs are skipped
unless --force is set (resumable).

Input formats accepted:
1. A list of strings:                ["text 1", "text 2", ...]
2. A list of objects with "text":    [{"id": "a", "text": "..."}, ...]
3. A dict mapping id → text:         {"a": "text 1", "b": "text 2"}

Usage:
    /mnt/modelall/ComfyUI/.venv/bin/python hymt2_batch.py vocab.json \\
        --out-dir /mnt/modelall/results/hymt2_vocab \\
        --target Japanese --prompt-lang en --ip 10.8.0.12 --limit 20
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
import tempfile
import time
from pathlib import Path

# Re-use the workflow + client from the single-shot CLI in this same dir.
sys.path.insert(0, str(Path(__file__).resolve().parent))
from hymt2_cli import (  # noqa: E402
    ComfyClient,
    build_workflow,
    extract_text_from_history,
)

_PREFERRED_TMP = "/mnt/modelall/tmp"
if os.path.isdir(_PREFERRED_TMP):
    tempfile.tempdir = _PREFERRED_TMP


# ---------------------------------------------------------------------------
# I/O helpers
# ---------------------------------------------------------------------------
_BAD_FN_RE = re.compile(r"[^\w\-\.\u4e00-\u9fff]+")  # keep word chars + CJK


def sanitize_filename(s: str, *, max_len: int = 80) -> str:
    """v2.0 §12 套路 C: CJK + special-character-safe filenames."""
    cleaned = _BAD_FN_RE.sub("_", s.strip()).strip("_") or "item"
    return cleaned[:max_len]


def load_items(path: Path) -> list[dict]:
    """Normalize the input JSON into a list of {"id": ..., "text": ...}."""
    data = json.loads(path.read_text(encoding="utf-8"))
    items: list[dict] = []
    if isinstance(data, list):
        for i, entry in enumerate(data):
            if isinstance(entry, str):
                items.append({"id": f"item_{i:06d}", "text": entry})
            elif isinstance(entry, dict):
                txt = entry.get("text") or entry.get("source") or entry.get("content")
                if not txt:
                    print(f"[hymt2/batch] skipping entry {i} — no text field",
                          file=sys.stderr)
                    continue
                eid = entry.get("id") or sanitize_filename(txt)[:40] or f"item_{i:06d}"
                items.append({"id": str(eid), "text": txt})
            else:
                print(f"[hymt2/batch] skipping entry {i} of type {type(entry)}",
                      file=sys.stderr)
    elif isinstance(data, dict):
        for k, v in data.items():
            if isinstance(v, str):
                items.append({"id": str(k), "text": v})
            elif isinstance(v, dict):
                txt = v.get("text") or v.get("source") or v.get("content")
                if txt:
                    items.append({"id": str(k), "text": txt})
    else:
        raise ValueError(f"unsupported input shape: {type(data)}")
    return items


def run_one(client: ComfyClient, *, source_text: str, args, force_reload: bool):
    wf = build_workflow(
        model_path=args.model_path,
        device=args.device,
        dtype=args.dtype,
        source_text=source_text,
        target_lang=args.target,
        prompt_lang=args.prompt_lang,
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=args.top_k,
        repetition_penalty=args.repetition_penalty,
        max_new_tokens=args.max_new_tokens,
        seed=args.seed,
        terminology=args.terminology_text,
        target_style=args.style,
        background_text=args.background_text,
        force_reload=force_reload,
    )
    history = client.run(wf, timeout_s=args.timeout)
    text = extract_text_from_history(history, "preview")
    if not text:
        text = extract_text_from_history(history, "translate")
    return text


def main():
    ap = argparse.ArgumentParser(description="Hy-MT2 batch translation.")
    ap.add_argument("input_json", type=Path,
                    help="Input JSON (list of strings, list of {id,text}, or dict).")

    ap.add_argument("--ip", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8188)
    ap.add_argument("--server", default=None)

    ap.add_argument("--model-path", default="tencent/Hy-MT2-1.8B")
    ap.add_argument("--gpu", type=int, default=0)
    ap.add_argument("--device", default=None)
    ap.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"])

    ap.add_argument("--target", default="English")
    ap.add_argument("--prompt-lang", default="en", choices=["en", "zh"])
    ap.add_argument("--temperature", type=float, default=0.7)
    ap.add_argument("--top-p", type=float, default=0.6)
    ap.add_argument("--top-k", type=int, default=20)
    ap.add_argument("--repetition-penalty", type=float, default=1.05)
    ap.add_argument("--max-new-tokens", type=int, default=4096)
    ap.add_argument("--seed", type=int, default=-1)

    ap.add_argument("--terminology", default="",
                    help="Path or inline glossary string applied to every item.")
    ap.add_argument("--style", default="")
    ap.add_argument("--background", default="")

    ap.add_argument("--out-dir", type=Path, required=True,
                    help="Output directory. One JSON per item.")
    ap.add_argument("--limit", type=int, default=0,
                    help="Process at most N items (0 = no limit).")
    ap.add_argument("--start", type=int, default=0,
                    help="Skip the first N items.")
    ap.add_argument("--force", action="store_true",
                    help="Overwrite existing output files.")
    ap.add_argument("--retries", type=int, default=2)
    ap.add_argument("--timeout", type=float, default=1800.0)
    ap.add_argument("--ws-recv-timeout", type=float, default=600.0)
    ap.add_argument("--dump-workflow", type=Path, default=None,
                    help="Dump the workflow of the FIRST item only (debug).")
    ap.add_argument("-v", "--verbose", action="store_true",
                    help="Print every received WebSocket message.")

    args = ap.parse_args()
    args.server = args.server or f"{args.ip}:{args.port}"
    args.device = args.device or f"cuda:{args.gpu}"

    # Resolve terminology / background once.
    def _resolve(value: str) -> str:
        if not value:
            return ""
        p = Path(value)
        if p.exists() and p.is_file():
            return p.read_text(encoding="utf-8")
        return value

    args.terminology_text = _resolve(args.terminology)
    args.background_text = _resolve(args.background)

    args.out_dir.mkdir(parents=True, exist_ok=True)

    items = load_items(args.input_json)
    print(f"[hymt2/batch] loaded {len(items)} items", file=sys.stderr)
    if args.start:
        items = items[args.start:]
    if args.limit:
        items = items[:args.limit]
    print(f"[hymt2/batch] will process {len(items)} items "
          f"(start={args.start} limit={args.limit})", file=sys.stderr)

    client = ComfyClient(server=args.server, ws_recv_timeout=args.ws_recv_timeout,
                         verbose=args.verbose)

    stats = {"ok": 0, "skip": 0, "fail": 0}
    t_run_start = time.time()
    first = True

    for idx, item in enumerate(items):
        eid = sanitize_filename(item["id"])
        out_path = args.out_dir / f"{eid}.json"
        if out_path.exists() and not args.force:
            stats["skip"] += 1
            continue

        if first and args.dump_workflow:
            wf_preview = build_workflow(
                model_path=args.model_path, device=args.device, dtype=args.dtype,
                source_text=item["text"], target_lang=args.target,
                prompt_lang=args.prompt_lang, temperature=args.temperature,
                top_p=args.top_p, top_k=args.top_k,
                repetition_penalty=args.repetition_penalty,
                max_new_tokens=args.max_new_tokens, seed=args.seed,
                terminology=args.terminology_text, target_style=args.style,
                background_text=args.background_text, force_reload=False,
            )
            args.dump_workflow.parent.mkdir(parents=True, exist_ok=True)
            args.dump_workflow.write_text(
                json.dumps(wf_preview, ensure_ascii=False, indent=2)
            )
            print(f"[hymt2/batch] dumped first workflow to {args.dump_workflow}",
                  file=sys.stderr)

        attempt = 0
        text = None
        last_err = None
        while attempt <= args.retries:
            try:
                text = run_one(client, source_text=item["text"], args=args,
                               force_reload=False)
                break
            except Exception as e:
                last_err = e
                attempt += 1
                print(f"[hymt2/batch] item={eid} attempt={attempt} "
                      f"failed: {e!r}", file=sys.stderr)
                time.sleep(2.0)

        first = False

        if text is None:
            stats["fail"] += 1
            print(f"[hymt2/batch] item={eid} GIVE UP after {args.retries} retries: "
                  f"{last_err!r}", file=sys.stderr)
            continue

        payload = {
            "id": item["id"],
            "source": item["text"],
            "translation": text,
            "target_lang": args.target,
            "prompt_lang": args.prompt_lang,
            "model_path": args.model_path,
        }
        out_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2),
                            encoding="utf-8")
        stats["ok"] += 1

        elapsed = time.time() - t_run_start
        done = stats["ok"] + stats["fail"]
        if done > 0:
            avg = elapsed / done
            eta = avg * (len(items) - idx - 1)
            print(f"[hymt2/batch] [{idx + 1}/{len(items)}] {eid} "
                  f"ok={stats['ok']} fail={stats['fail']} skip={stats['skip']} "
                  f"avg={avg:.1f}s/item eta={eta / 60:.1f}min",
                  file=sys.stderr)

    print(f"[hymt2/batch] DONE ok={stats['ok']} fail={stats['fail']} "
          f"skip={stats['skip']}", file=sys.stderr)


if __name__ == "__main__":
    main()