#!/usr/bin/env python3
"""hymt2_cli.py — single-shot Hy-MT2 translation via ComfyUI.

Usage:
    python hymt2_cli.py "今天天气真好。" --target English --ip 10.8.0.12
    python hymt2_cli.py - --target 日语 --prompt-lang zh --ip 10.8.0.12  # stdin
    python hymt2_cli.py path/to/text.txt --target French --ip 10.8.0.12
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import tempfile
import time
import urllib.parse
import urllib.request
import uuid
from pathlib import Path

import websocket  # pip install websocket-client


# Use /mnt/modelall/tmp if available (v2.0 §1: /tmp is only 10 GB on the server).
_PREFERRED_TMP = "/mnt/modelall/tmp"
if os.path.isdir(_PREFERRED_TMP):
    tempfile.tempdir = _PREFERRED_TMP


# ----------------------------------------------------------------------------
# WebSocket completion signals
# ----------------------------------------------------------------------------
# Modern ComfyUI (the version Hua deployed in 2026-05) emits a different set of
# WS messages than the older docs describe. From actual /history capture:
#   prompt_id: ...
#   status messages:
#     - execution_start
#     - execution_cached    (some nodes hit the cache)
#     - execution_success   (whole prompt finished)
#
# Crucially, in fully-cached or very fast paths, the older `executing` /
# `node=null` "queue empty" message is NOT emitted at all, so a client that
# only listens for that signal will block on ws.recv() forever.
#
# v0.1.1 handles all of:
#   - execution_success  -> our prompt completed cleanly
#   - execution_error    -> our prompt failed
#   - executing + node=None -> legacy completion signal (still emitted on the
#                             non-cached path)
#   - executed           -> single-node output (informational, keep listening)
#   - execution_cached   -> informational only
#   - progress           -> progress bar

_COMPLETION_TYPES = {"execution_success"}
_ERROR_TYPES = {"execution_error"}


class ComfyClient:
    """Minimal ComfyUI client.

    v0.1.2: opens WS BEFORE submit + polls /history on status messages to
            avoid the cached-prompt race that would block on ws.recv().
    v0.1.1: handles modern execution_success WS messages.
    """

    def __init__(self, server="127.0.0.1:8188", *,
                 ws_recv_timeout=600.0, connect_timeout=30.0, verbose=False):
        self.server = server
        self.client_id = str(uuid.uuid4())
        self.ws_recv_timeout = ws_recv_timeout
        self.connect_timeout = connect_timeout
        self.verbose = verbose

    def submit(self, workflow):
        payload = json.dumps({"prompt": workflow, "client_id": self.client_id}).encode()
        req = urllib.request.Request(
            f"http://{self.server}/prompt",
            data=payload,
            headers={"Content-Type": "application/json"},
        )
        with urllib.request.urlopen(req, timeout=self.connect_timeout) as r:
            data = json.loads(r.read())
        return data["prompt_id"]

    def history(self, prompt_id):
        url = f"http://{self.server}/history/{urllib.parse.quote(prompt_id)}"
        with urllib.request.urlopen(url, timeout=self.connect_timeout) as r:
            data = json.loads(r.read())
        return data.get(prompt_id, {})

    def _log(self, msg):
        if self.verbose:
            print(f"[hymt2/cli] {msg}", file=sys.stderr, flush=True)

    def run(self, workflow, *, timeout_s=1800):
        # v0.1.2: open WS BEFORE submit to avoid the cached-prompt race.
        # If we submit first then connect, the server may broadcast all
        # execution_start / execution_success events before our WS is up,
        # and we'd block on ws.recv() forever waiting for events that already
        # happened. Opening WS first means the server's broadcast lands in our
        # socket buffer.
        ws = websocket.WebSocket()
        ws.connect(
            f"ws://{self.server}/ws?clientId={self.client_id}",
            timeout=self.connect_timeout,
        )
        # ⭐ v2.0 §10: ws.connect(timeout) is connect-only. Set the recv timeout
        # so cold load (model download / first-load) doesn't trip the socket.
        ws.settimeout(self.ws_recv_timeout)

        try:
            prompt_id = self.submit(workflow)
            self._log(f"submitted prompt_id={prompt_id}")

            # v0.1.2: Even with WS-first, an instant cache hit can complete
            # during the submit() POST round trip. Poll /history once before
            # entering the WS loop.
            h0 = self.history(prompt_id)
            if h0 and h0.get("status", {}).get("completed"):
                self._log("prompt already completed at submit time (cache hit)")
                return h0

            deadline = time.time() + timeout_s
            completed = False
            last_exec_err = None

            while not completed:
                if time.time() > deadline:
                    raise TimeoutError(f"workflow exceeded {timeout_s}s overall")

                try:
                    raw = ws.recv()
                except (websocket.WebSocketTimeoutException, TimeoutError):
                    # Idle timeout — poll /history as fallback. Server may have
                    # completed and we just missed the WS signal.
                    self._log("ws recv idle timeout, polling /history")
                    h = self.history(prompt_id)
                    if h and h.get("status", {}).get("completed"):
                        self._log("fallback /history poll found completion")
                        return h
                    continue
                except websocket.WebSocketConnectionClosedException:
                    # Server closed the socket. Fall through to /history.
                    self._log("ws connection closed, polling /history")
                    break

                if isinstance(raw, bytes):
                    # Binary frames are progress images we don't care about.
                    continue
                try:
                    msg = json.loads(raw)
                except Exception:
                    continue

                kind = msg.get("type")
                data = msg.get("data", {}) or {}
                msg_pid = data.get("prompt_id")

                if self.verbose:
                    s = json.dumps(data, ensure_ascii=False, default=str)
                    if len(s) > 200:
                        s = s[:200] + "...<truncated>"
                    self._log(f"ws msg type={kind} data={s}")

                # --- Completion: execution_success for our prompt ---
                if kind in _COMPLETION_TYPES and msg_pid == prompt_id:
                    completed = True
                    break

                # --- Error: execution_error for our prompt ---
                if kind in _ERROR_TYPES and msg_pid == prompt_id:
                    last_exec_err = data
                    break

                # --- Legacy completion signal: executing/node=null ---
                if kind == "executing" and data.get("node") is None:
                    # Some older versions don't set prompt_id; accept either
                    # match or missing.
                    if msg_pid is None or msg_pid == prompt_id:
                        completed = True
                        break

                # --- Progress bar ---
                if kind == "progress":
                    val = data.get("value")
                    mx = data.get("max")
                    if val is not None and mx:
                        print(f"  progress {val}/{mx}",
                              file=sys.stderr, flush=True)

                # v0.1.2: status with queue_remaining=0 is a strong hint that
                # the server already finished. Cross-check /history.
                if kind == "status":
                    qrem = (data.get("status", {})
                                .get("exec_info", {})
                                .get("queue_remaining"))
                    if qrem == 0:
                        h = self.history(prompt_id)
                        if h and h.get("status", {}).get("completed"):
                            self._log("status/queue_remaining=0 + history "
                                      "shows completed -> race detected")
                            return h

                # Other types (execution_start, execution_cached, executed)
                # are informational; we keep listening.
        finally:
            try:
                ws.close()
            except Exception:
                pass

        if last_exec_err:
            raise RuntimeError(
                f"ComfyUI execution_error: "
                f"{json.dumps(last_exec_err, ensure_ascii=False)[:1000]}"
            )

        # ALWAYS fetch /history at the end. Outputs are only guaranteed to be
        # in /history; the WebSocket success message is just a signal.
        h = self.history(prompt_id)
        if self.verbose:
            status = h.get("status", {})
            self._log(f"final /history: completed={status.get('completed')} "
                      f"status_str={status.get('status_str')}")
        return h


def extract_text_from_history(history, node_id) -> str:
    """v2.0 §5.3 — peel STRING outputs out of /history."""
    outs = history.get("outputs", {})
    if node_id not in outs:
        return ""
    no = outs[node_id]
    if "text" in no:
        v = no["text"]
        return "\n".join(str(x) for x in v) if isinstance(v, list) else str(v)
    parts = []
    for v in no.values():
        if isinstance(v, list):
            parts.extend(str(x) for x in v)
    return "\n".join(parts)


def build_workflow(*, model_path, device, dtype, source_text, target_lang,
                   prompt_lang, temperature, top_p, top_k, repetition_penalty,
                   max_new_tokens, seed, terminology, target_style,
                   background_text, force_reload):
    """Return the API-format ComfyUI workflow dict (v2.0 §5.2 style)."""
    return {
        "loader": {
            "class_type": "HyMT2Loader",
            "inputs": {
                "model_path": model_path,
                "device": device,
                "dtype": dtype,
                "force_reload": force_reload,
                "evict_others": True,
            },
        },
        "translate": {
            "class_type": "HyMT2Translate",
            "inputs": {
                "model": ["loader", 0],
                "source_text": source_text,
                "target_lang": target_lang,
                "prompt_lang": prompt_lang,
                "temperature": temperature,
                "top_p": top_p,
                "top_k": top_k,
                "repetition_penalty": repetition_penalty,
                "max_new_tokens": max_new_tokens,
                "seed": seed,
                "terminology": terminology,
                "target_style": target_style,
                "background_text": background_text,
            },
        },
        "preview": {
            "class_type": "PreviewAny",
            "inputs": {"source": ["translate", 0]},
        },
    }


def read_source(src_arg: str) -> str:
    """Resolve CLI arg into actual text:
       - '-'  → stdin
       - a path that exists → file contents
       - else → treat as inline text
    """
    if src_arg == "-":
        return sys.stdin.read()
    p = Path(src_arg)
    if p.exists() and p.is_file():
        return p.read_text(encoding="utf-8")
    return src_arg


def main():
    ap = argparse.ArgumentParser(description="Hy-MT2 translation via ComfyUI.")
    ap.add_argument("source",
                    help="Source text. Use '-' for stdin, or a file path, "
                         "or inline string.")

    # Server (v2.0 §6)
    ap.add_argument("--ip", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8188)
    ap.add_argument("--server", default=None,
                    help="Full host:port, overrides --ip/--port.")

    # Model selection
    ap.add_argument("--model-path", default="tencent/Hy-MT2-1.8B",
                    help="HF repo id or absolute local dir.")
    ap.add_argument("--gpu", type=int, default=0,
                    help="CUDA device index → cuda:N")
    ap.add_argument("--device", default=None,
                    help="Override --gpu with an explicit device string.")
    ap.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"])

    # Translation params
    ap.add_argument("--target", default="English",
                    help="Target language NAME. Use English names with "
                         "--prompt-lang en, Chinese names with --prompt-lang zh.")
    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)

    # Optional template extras
    ap.add_argument("--terminology", default="",
                    help="Path to a glossary file (multi-line, 'src=>tgt' format) "
                         "or the inline glossary string.")
    ap.add_argument("--style", default="", help="Target style descriptor.")
    ap.add_argument("--background", default="",
                    help="Background paragraph or path to a file.")

    # Output / debug (v2.0 §6 standard set)
    ap.add_argument("--out", type=Path, default=None,
                    help="Write the translation to this file. "
                         "If not set, print to stdout.")
    ap.add_argument("--dump-workflow", type=Path, default=None,
                    help="Dump the generated API workflow JSON for UI loading.")
    ap.add_argument("--timeout", type=float, default=1800.0)
    ap.add_argument("--ws-recv-timeout", type=float, default=600.0)
    ap.add_argument("--force-reload", action="store_true",
                    help="Skip the loader cache (force re-load).")
    ap.add_argument("-v", "--verbose", action="store_true",
                    help="Print every received WebSocket message and the "
                         "final /history status.")

    args = ap.parse_args()

    server = args.server or f"{args.ip}:{args.port}"
    device = args.device or f"cuda:{args.gpu}"

    source_text = read_source(args.source)

    def _resolve_optional_text(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

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

    wf = build_workflow(
        model_path=args.model_path,
        device=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=terminology_text,
        target_style=args.style,
        background_text=background_text,
        force_reload=args.force_reload,
    )

    if args.dump_workflow:
        args.dump_workflow.parent.mkdir(parents=True, exist_ok=True)
        args.dump_workflow.write_text(json.dumps(wf, ensure_ascii=False, indent=2))
        print(f"workflow dumped to {args.dump_workflow}", file=sys.stderr)

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

    t0 = time.time()
    history = client.run(wf, timeout_s=args.timeout)
    dt = time.time() - t0

    text = extract_text_from_history(history, "preview")
    if not text:
        text = extract_text_from_history(history, "translate")

    print(f"[hymt2/cli] done in {dt:.1f}s", file=sys.stderr)

    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(text, encoding="utf-8")
        print(f"[hymt2/cli] wrote {len(text)} chars to {args.out}", file=sys.stderr)
    else:
        print(text)


if __name__ == "__main__":
    main()