#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "requests",
#   "websocket-client",
# ]
# ///
"""mimo_asr_cli.py — drive ComfyUI-MiMoASR from the command line.

Examples
--------
  # 本地音频 → 自动语言检测
  uv run python cli/mimo_asr_cli.py audio.wav --ip 10.8.0.12

  # 显式指定中文
  uv run python cli/mimo_asr_cli.py speech.flac --language Chinese --ip 10.8.0.12

  # 自定义模型路径(默认是 ./models/MiMo-V2.5-ASR + ./models/MiMo-Audio-Tokenizer
  # 相对于 ComfyUI 启动目录)
  uv run python cli/mimo_asr_cli.py speech.wav \
      --model-path /mnt/models/MiMo-V2.5-ASR \
      --tokenizer-path /mnt/models/MiMo-Audio-Tokenizer \
      --repo-path /mnt/repos/MiMo-V2.5-ASR \
      --ip 10.8.0.12

  # 调试:dump 生成的 workflow 不实际执行
  uv run python cli/mimo_asr_cli.py speech.wav --dump-workflow /tmp/wf.json --no-run
"""

from __future__ import annotations

import argparse
import json
import mimetypes
import re
import sys
import time
import urllib.parse
import uuid
from pathlib import Path

import requests
import websocket  # websocket-client

LANGUAGES = ["Auto", "Chinese", "English"]


# ---------------------------------------------------------------------------
# ComfyClient — minimal HTTP + WS wrapper around a ComfyUI server
# ---------------------------------------------------------------------------
class ComfyClient:
    def __init__(self, server: str = "127.0.0.1:8188") -> None:
        self.server = server
        self.http = f"http://{server}"
        self.ws = f"ws://{server}/ws"
        self.client_id = str(uuid.uuid4())

    def upload_image(self, file_path: Path, *, overwrite: bool = True) -> str:
        """Upload any file via POST /upload/image. Works for audio too."""
        mime, _ = mimetypes.guess_type(str(file_path))
        with open(file_path, "rb") as fh:
            files = {"image": (file_path.name, fh, mime or "application/octet-stream")}
            data = {"overwrite": "true" if overwrite else "false"}
            r = requests.post(f"{self.http}/upload/image", files=files, data=data, timeout=60)
        r.raise_for_status()
        payload = r.json()
        # ComfyUI returns {"name": "<server-side filename>", ...}
        return payload["name"]

    def object_info(self) -> dict:
        r = requests.get(f"{self.http}/object_info", timeout=15)
        r.raise_for_status()
        return r.json()

    def submit(self, workflow: dict) -> str:
        r = requests.post(
            f"{self.http}/prompt",
            json={"prompt": workflow, "client_id": self.client_id},
            timeout=30,
        )
        r.raise_for_status()
        payload = r.json()
        if "prompt_id" not in payload:
            raise RuntimeError(f"Submit failed: {payload}")
        return payload["prompt_id"]

    def history(self, prompt_id: str) -> dict:
        r = requests.get(f"{self.http}/history/{prompt_id}", timeout=15)
        r.raise_for_status()
        data = r.json()
        return data.get(prompt_id, {})

    def view(self, filename: str, subfolder: str, type_: str) -> bytes:
        params = urllib.parse.urlencode(
            {"filename": filename, "subfolder": subfolder, "type": type_}
        )
        r = requests.get(f"{self.http}/view?{params}", timeout=60)
        r.raise_for_status()
        return r.content

    def run(self, workflow: dict, *, timeout_s: float = 600.0) -> dict:
        """Submit a workflow and stream progress until completion."""
        prompt_id = self.submit(workflow)
        print(f"  prompt_id = {prompt_id}")

        ws = websocket.create_connection(f"{self.ws}?clientId={self.client_id}", timeout=30)
        ws.settimeout(timeout_s)
        deadline = time.time() + timeout_s
        last_node = None
        try:
            while True:
                if time.time() > deadline:
                    raise TimeoutError(f"Workflow exceeded {timeout_s}s")
                msg = ws.recv()
                if isinstance(msg, (bytes, bytearray)):
                    continue
                try:
                    payload = json.loads(msg)
                except json.JSONDecodeError:
                    continue
                t = payload.get("type")
                d = payload.get("data", {})
                if d.get("prompt_id") and d.get("prompt_id") != prompt_id:
                    continue
                if t == "progress":
                    cur, tot = d.get("value", 0), d.get("max", 1)
                    sys.stdout.write(f"\r  ▸ {cur:>4}/{tot}")
                    sys.stdout.flush()
                elif t == "executing":
                    node = d.get("node")
                    if node is None:
                        # execution finished for this prompt
                        sys.stdout.write("\n")
                        break
                    if node != last_node:
                        sys.stdout.write(f"\r  ▸ node {str(node):<20}\n")
                        sys.stdout.flush()
                        last_node = node
                elif t == "execution_error":
                    sys.stdout.write("\n")
                    raise RuntimeError(
                        f"Workflow error in node {d.get('node_id')}: "
                        f"{d.get('exception_message')}"
                    )
        finally:
            try:
                ws.close()
            except Exception:  # noqa: BLE001
                pass
        return self.history(prompt_id)


# ---------------------------------------------------------------------------
# Workflow construction
# ---------------------------------------------------------------------------
def build_workflow(
    *,
    audio_filename: str,
    model_path: str,
    tokenizer_path: str,
    repo_path: str,
    language: str,
) -> dict:
    """Return the API-format workflow dict."""
    return {
        "loader": {
            "class_type": "MiMoASRLoader",
            "inputs": {
                "model_path": model_path,
                "tokenizer_path": tokenizer_path,
                "repo_path": repo_path,
                "force_reload": False,
            },
        },
        "load_audio": {
            "class_type": "LoadAudio",
            "inputs": {
                "audio": audio_filename,
            },
        },
        "asr": {
            "class_type": "MiMoASRTranscribe",
            "inputs": {
                "model": ["loader", 0],
                "audio": ["load_audio", 0],
                "language": language,
            },
        },
        # Capture STRING output so we can read it back from /history.
        "preview": {
            "class_type": "PreviewAny",
            "inputs": {
                "source": ["asr", 0],
            },
        },
    }


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def extract_text_from_history(history: dict, node_id: str = "preview") -> str:
    """Pull the text payload PreviewAny stored under outputs[node_id]."""
    outs = history.get("outputs", {})
    if node_id not in outs:
        # Fall back to the asr node itself if PreviewAny was missing.
        for nid, node_out in outs.items():
            if "text" in node_out:
                node_id = nid
                break
        else:
            return ""
    node_out = outs[node_id]
    if "text" in node_out:
        v = node_out["text"]
        return "\n".join(str(x) for x in v) if isinstance(v, list) else str(v)
    parts: list[str] = []
    for v in node_out.values():
        if isinstance(v, list):
            parts.extend(str(x) for x in v)
    return "\n".join(parts)


_SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._-]+")


def sanitize_filename(name: str) -> str:
    sanitized = _SAFE_NAME_RE.sub("_", name).strip("._-")
    return sanitized or "audio"


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------
def main() -> int:
    ap = argparse.ArgumentParser(description="ComfyUI-MiMoASR CLI")
    ap.add_argument("audio", type=Path, help="Local audio file to transcribe (wav/mp3/flac/m4a/...)")

    # Server.
    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")

    # Output.
    ap.add_argument("--out", type=Path, default=None,
                    help="Write transcript to file (default: stdout only).")

    # MiMo paths.
    ap.add_argument("--model-path", default="./models/MiMo-V2.5-ASR",
                    help="Path to MiMo-V2.5-ASR weights (relative to ComfyUI cwd or absolute).")
    ap.add_argument("--tokenizer-path", default="./models/MiMo-Audio-Tokenizer",
                    help="Path to MiMo-Audio-Tokenizer.")
    ap.add_argument("--repo-path", default="",
                    help="Path to local clone of XiaomiMiMo/MiMo-V2.5-ASR. "
                         "Leave blank to use the MIMO_ASR_REPO env var on the server.")

    # Generation.
    ap.add_argument("--language", choices=LANGUAGES, default="Auto")

    # Debug.
    ap.add_argument("--dump-workflow", type=Path, default=None,
                    help="Save the generated workflow JSON for inspection.")
    ap.add_argument("--no-run", action="store_true",
                    help="Skip submission (use with --dump-workflow to just produce the JSON).")
    ap.add_argument("--timeout", type=float, default=600.0)

    args = ap.parse_args()

    server = args.server or f"{args.ip}:{args.port}"
    if not args.audio.is_file():
        print(f"Error: audio file not found: {args.audio}", file=sys.stderr)
        return 2

    client = ComfyClient(server=server)

    # 1) Upload audio.
    if not args.no_run:
        print(f"▸ Uploading {args.audio.name} to {server} ...")
        uploaded_name = client.upload_image(args.audio)
        print(f"  uploaded as: {uploaded_name}")
    else:
        uploaded_name = sanitize_filename(args.audio.name)

    # 2) Build workflow.
    workflow = build_workflow(
        audio_filename=uploaded_name,
        model_path=args.model_path,
        tokenizer_path=args.tokenizer_path,
        repo_path=args.repo_path,
        language=args.language,
    )

    if args.dump_workflow:
        args.dump_workflow.parent.mkdir(parents=True, exist_ok=True)
        args.dump_workflow.write_text(
            json.dumps(workflow, indent=2, ensure_ascii=False), encoding="utf-8"
        )
        print(f"▸ Workflow JSON written to {args.dump_workflow}")

    if args.no_run:
        print("▸ --no-run set, exiting without submission.")
        return 0

    # 3) Run.
    print(f"▸ Submitting workflow (language={args.language!r}) ...")
    t0 = time.time()
    history = client.run(workflow, timeout_s=args.timeout)
    elapsed = time.time() - t0
    print(f"▸ Workflow finished in {elapsed:.2f}s")

    # 4) Extract text.
    text = extract_text_from_history(history, node_id="preview")
    if not text:
        print("Warning: no text output found in workflow history.", file=sys.stderr)
        print(json.dumps(history.get("outputs", {}), indent=2, ensure_ascii=False)[:2000],
              file=sys.stderr)
        return 1

    print("\n----- Transcript -----")
    print(text)
    print("----------------------")

    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(text + "\n", encoding="utf-8")
        print(f"▸ Saved to {args.out}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())