#!/usr/bin/env python3
"""
video_to_text_cli.py - Submit a video to ComfyUI-VideoToText + ComfyUI-MiMoASR and dump the JSON transcript.

Pipeline: VTT_LoadVideo -> VTT_ExtractAudio -> VTT_AudioFromPath
           -> MiMoASRLoader + MiMoASRTranscribe
           -> VTT_BuildTranscriptJSON -> VTT_SaveTranscript

Usage:
    uv run python video_to_text_cli.py /path/to/video.mp4 \
        --ip 10.8.0.12 \
        --model-path /big_disk/models/MiMo-V2.5-ASR \
        --tokenizer-path /big_disk/models/MiMo-Audio-Tokenizer \
        --repo-path /opt/repos/MiMo-V2.5-ASR \
        --language Auto

Dependencies:
    uv pip install --system websocket-client requests
"""
from __future__ import annotations

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

import requests
import websocket


# ----------------------------------------------------------------------
# ComfyClient - same shape as the rest of the aadebuger CLI fleet.
# ----------------------------------------------------------------------
class ComfyClient:
    def __init__(self, server: str = "127.0.0.1:8188"):
        self.server = server
        self.client_id = str(uuid.uuid4())
        self.base = f"http://{server}"

    def upload_file(self, file_path: Path, *, overwrite: bool = True) -> str:
        """POST /upload/image - endpoint is named 'image' but accepts any file."""
        with open(file_path, "rb") as f:
            files = {"image": (file_path.name, f, "application/octet-stream")}
            data = {"overwrite": "true" if overwrite else "false"}
            r = requests.post(f"{self.base}/upload/image", files=files, data=data, timeout=600)
            r.raise_for_status()
            return r.json().get("name", file_path.name)

    def submit(self, workflow: dict) -> str:
        payload = {"prompt": workflow, "client_id": self.client_id}
        r = requests.post(f"{self.base}/prompt", json=payload, timeout=30)
        r.raise_for_status()
        return r.json()["prompt_id"]

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

    def view(self, filename: str, subfolder: str = "", type_: str = "output") -> bytes:
        params = {"filename": filename, "subfolder": subfolder, "type": type_}
        r = requests.get(f"{self.base}/view", params=params, timeout=300)
        r.raise_for_status()
        return r.content

    def run(self, workflow: dict, *, timeout_s: float = 1800) -> dict:
        ws = websocket.WebSocket()
        ws.connect(f"ws://{self.server}/ws?clientId={self.client_id}")
        prompt_id = self.submit(workflow)
        start = time.time()
        last_node = None
        try:
            while True:
                if time.time() - start > timeout_s:
                    raise TimeoutError(f"Timed out after {timeout_s:.0f}s")
                msg = ws.recv()
                if not isinstance(msg, str):
                    continue
                m = json.loads(msg)
                t = m.get("type")
                d = m.get("data", {})
                if t == "executing":
                    if d.get("prompt_id") != prompt_id:
                        continue
                    nid = d.get("node")
                    if nid is None:
                        print()
                        break
                    if nid != last_node:
                        print(f"\r  > node {str(nid):<28}", end="", flush=True)
                        last_node = nid
                elif t == "progress":
                    if d.get("prompt_id") != prompt_id:
                        continue
                    cur, tot = d.get("value", 0), d.get("max", 1)
                    print(f"\r  > {str(last_node or '?'):<14} {cur:>4}/{tot}", end="", flush=True)
                elif t == "execution_error":
                    raise RuntimeError(f"Node error: {json.dumps(d, ensure_ascii=False)[:600]}")
        finally:
            ws.close()
        return self.history(prompt_id)


# ----------------------------------------------------------------------
# build_workflow - API format; semantic node ids.
# ----------------------------------------------------------------------
def build_workflow(
    *,
    video_filename: str,
    model_path: str,
    tokenizer_path: str,
    repo_path: str,
    language: str,
    sample_rate: int,
    save_prefix: str,
    force_reload: bool = False,
) -> dict:
    return {
        "loader_video": {
            "class_type": "VTT_LoadVideo",
            "inputs": {
                "video": video_filename,
                "video_path_override": "",
            },
        },
        "extract_audio": {
            "class_type": "VTT_ExtractAudio",
            "inputs": {
                "video_path": ["loader_video", 0],
                "sample_rate": sample_rate,
                "mono": True,
            },
        },
        "audio_in": {
            "class_type": "VTT_AudioFromPath",
            "inputs": {
                "audio_path": ["extract_audio", 0],
                "target_sample_rate": 0,
                "force_mono": False,
            },
        },
        "mimo_loader": {
            "class_type": "MiMoASRLoader",
            "inputs": {
                "model_path": model_path,
                "tokenizer_path": tokenizer_path,
                "repo_path": repo_path,
                "force_reload": force_reload,
            },
        },
        "mimo_asr": {
            "class_type": "MiMoASRTranscribe",
            "inputs": {
                "model": ["mimo_loader", 0],
                "audio": ["audio_in", 0],
                "language": language,
            },
        },
        "build_json": {
            "class_type": "VTT_BuildTranscriptJSON",
            "inputs": {
                "text": ["mimo_asr", 0],
                "audio_path": ["extract_audio", 0],
                "duration_sec": ["audio_in", 1],
                "model_name": "MiMo-V2.5-ASR",
                "tokenizer_name": "MiMo-Audio-Tokenizer",
                "language": language,
                "source_video_path": ["loader_video", 0],
            },
        },
        "save": {
            "class_type": "VTT_SaveTranscript",
            "inputs": {
                "transcript_json": ["build_json", 0],
                "filename_prefix": save_prefix,
                "include_timestamp": True,
                "subfolder": "",
            },
        },
        "preview_text": {
            "class_type": "PreviewAny",
            "inputs": {"source": ["mimo_asr", 0]},
        },
        "preview_json": {
            "class_type": "PreviewAny",
            "inputs": {"source": ["build_json", 0]},
        },
    }


# ----------------------------------------------------------------------
# extract_text_from_history - pull STRING output from PreviewAny.
# ----------------------------------------------------------------------
def extract_text_from_history(history: dict, node_id: str) -> str:
    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 is_url(s: str) -> bool:
    try:
        u = urlparse(s)
        return u.scheme in {"http", "https"} and bool(u.netloc)
    except Exception:
        return False


# ----------------------------------------------------------------------
# main
# ----------------------------------------------------------------------
def main() -> int:
    ap = argparse.ArgumentParser(
        description="Submit a video to ComfyUI-VideoToText + ComfyUI-MiMoASR and dump the JSON transcript."
    )
    ap.add_argument(
        "video",
        help="Local video path (will be uploaded) OR filename already in ComfyUI/input.",
    )

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

    # MiMo paths (required on first run; can hardcode in a wrapper script)
    ap.add_argument(
        "--model-path", required=True,
        help="Absolute path to MiMo-V2.5-ASR on the server (e.g. /big_disk/models/MiMo-V2.5-ASR)."
    )
    ap.add_argument(
        "--tokenizer-path", required=True,
        help="Absolute path to MiMo-Audio-Tokenizer on the server."
    )
    ap.add_argument(
        "--repo-path", default="",
        help="Absolute path to MiMo-V2.5-ASR source repo. Leave blank if MIMO_ASR_REPO is set."
    )
    ap.add_argument("--force-reload", action="store_true",
                    help="Reload the model even if cached.")

    # Inference
    ap.add_argument("--language", default="Auto",
                    choices=["Auto", "Chinese", "English"])
    ap.add_argument("--sample-rate", type=int, default=16000)
    ap.add_argument("--save-prefix", default="transcript")

    # Output / debug
    ap.add_argument("--out", type=Path, default=Path("./out"))
    ap.add_argument("--dump-workflow", type=Path, default=None,
                    help="Dump generated workflow JSON for debugging.")
    ap.add_argument("--timeout", type=float, default=1800)

    args = ap.parse_args()

    server = args.server or f"{args.ip}:{args.port}"
    client = ComfyClient(server=server)

    # Resolve input
    if is_url(args.video):
        print("[!] URL input not supported. Download first or upload to ComfyUI/input.",
              file=sys.stderr)
        return 2

    src = Path(args.video).expanduser()
    if src.exists() and src.is_file():
        size_mb = src.stat().st_size / 1e6
        print(f"[+] Uploading {src.name} ({size_mb:.1f} MB) -> {server} ...")
        uploaded_name = client.upload_file(src)
        print(f"    -> ComfyUI/input/{uploaded_name}")
    else:
        uploaded_name = args.video
        print(f"[+] Using existing ComfyUI/input/{uploaded_name}")

    workflow = build_workflow(
        video_filename=uploaded_name,
        model_path=args.model_path,
        tokenizer_path=args.tokenizer_path,
        repo_path=args.repo_path,
        language=args.language,
        sample_rate=args.sample_rate,
        save_prefix=args.save_prefix,
        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(workflow, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        print(f"[+] Dumped workflow -> {args.dump_workflow}")

    print(f"[+] Submitting (model=MiMo-V2.5-ASR, language={args.language}) ...")
    history = client.run(workflow, timeout_s=args.timeout)

    transcript = extract_text_from_history(history, "preview_json")
    text_only = extract_text_from_history(history, "preview_text")

    if not transcript:
        print("[!] No transcript output. Dumping head of history:", file=sys.stderr)
        print(json.dumps(history, ensure_ascii=False, indent=2)[:2000], file=sys.stderr)
        return 1

    args.out.mkdir(parents=True, exist_ok=True)
    stem = src.stem if src.exists() else Path(args.video).stem
    out_path = args.out / f"{stem}.json"
    out_path.write_text(transcript, encoding="utf-8")

    # Summary
    try:
        data = json.loads(transcript)
        dur = data.get("source", {}).get("duration_sec", 0.0)
        chars = len(data.get("text", ""))
        print(f"[v] duration {dur:.1f}s | {chars} chars | language={args.language}")
    except Exception:
        pass

    print(f"[v] Saved -> {out_path}")
    if text_only:
        preview = text_only[:200].replace("\n", " ")
        print(f"[v] Preview: {preview}{'...' if len(text_only) > 200 else ''}")
    return 0


if __name__ == "__main__":
    sys.exit(main())