ComfyUI-ShotSplitter/server_hooks.py

"""ComfyUI server hooks for Oz ShotSplitter.

Routes:
  GET  /oz_shotsplitter/active_clients  — debug list of active websocket clients
  POST /oz_shotsplitter/download_url    — yt-dlp a public URL into ComfyUI input/
"""

from __future__ import annotations

import asyncio
import hashlib
import logging
import os
import re

logger = logging.getLogger(__name__)

try:
    from server import PromptServer  # type: ignore
except Exception as e:  # pragma: no cover
    PromptServer = None  # type: ignore
    logger.warning("Oz ShotSplitter: PromptServer unavailable (%s)", e)

DOWNLOAD_TIMEOUT_SEC = 600
DOWNLOAD_MAX_FILESIZE = "500M"
VIDEO_EXTS = ("mp4", "mkv", "webm", "mov", "m4v")
_INPUT_DIR_CACHE: str | None = None


def _resolve_input_dir() -> str:
    global _INPUT_DIR_CACHE
    if _INPUT_DIR_CACHE:
        return _INPUT_DIR_CACHE
    try:
        import folder_paths  # type: ignore
        d = folder_paths.get_input_directory()
    except Exception:
        d = "/workspace/ComfyUI/input"
    os.makedirs(d, exist_ok=True)
    _INPUT_DIR_CACHE = d
    return d


def _find_cached(input_dir: str, stem: str) -> str | None:
    for ext in VIDEO_EXTS:
        p = os.path.join(input_dir, f"{stem}.{ext}")
        if os.path.exists(p) and os.path.getsize(p) > 0:
            return p
    return None


def register_routes() -> None:
    if PromptServer is None:
        return
    try:
        server = PromptServer.instance
    except Exception as e:
        logger.warning("Oz ShotSplitter: no PromptServer.instance (%s)", e)
        return

    from aiohttp import web

    @server.routes.get("/oz_shotsplitter/active_clients")
    async def _list_active_clients(_request):
        sids = list(getattr(server, "sockets", {}).keys())
        return web.json_response({"client_ids": sids, "count": len(sids)})

    @server.routes.post("/oz_shotsplitter/download_url")
    async def _download_url(request):
        try:
            data = await request.json()
        except Exception:
            return web.json_response({"ok": False, "error": "invalid JSON body"}, status=400)

        url = (data.get("url") or "").strip()
        if not url or not re.match(r"^https?://", url):
            return web.json_response({"ok": False, "error": "invalid URL"}, status=400)

        input_dir = _resolve_input_dir()
        url_hash = hashlib.md5(url.encode("utf-8")).hexdigest()[:12]
        stem = f"yt_{url_hash}"

        cached = _find_cached(input_dir, stem)
        if cached:
            return web.json_response({
                "ok": True,
                "filename": os.path.basename(cached),
                "cached": True,
            })

        out_template = os.path.join(input_dir, f"{stem}.%(ext)s")
        cmd = [
            "yt-dlp",
            url,
            "-o", out_template,
            "--no-playlist",
            "--max-filesize", DOWNLOAD_MAX_FILESIZE,
            "-f", "mp4/best[ext=mp4]/bestvideo*+bestaudio/best",
            "--merge-output-format", "mp4",
            "--no-warnings",
            "--no-progress",
            "--quiet",
        ]
        logger.info("Oz ShotSplitter: yt-dlp %s -> %s", url, stem)

        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
        except FileNotFoundError:
            return web.json_response({
                "ok": False,
                "error": "yt-dlp not installed. Run: /venv/main/bin/pip install yt-dlp",
            }, status=500)
        except Exception as e:
            return web.json_response({"ok": False, "error": f"spawn failed: {e}"}, status=500)

        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=DOWNLOAD_TIMEOUT_SEC
            )
        except asyncio.TimeoutError:
            try:
                proc.kill()
                await proc.wait()
            except Exception:
                pass
            return web.json_response({
                "ok": False,
                "error": f"yt-dlp timed out after {DOWNLOAD_TIMEOUT_SEC}s",
            }, status=504)

        if proc.returncode != 0:
            err_tail = (stderr or b"").decode("utf-8", errors="replace").strip()
            if len(err_tail) > 600:
                err_tail = err_tail[-600:]
            return web.json_response({
                "ok": False,
                "error": f"yt-dlp exit {proc.returncode}: {err_tail or 'no stderr'}",
            }, status=500)

        saved = _find_cached(input_dir, stem)
        if not saved:
            return web.json_response({
                "ok": False,
                "error": "yt-dlp succeeded but output file missing",
            }, status=500)

        return web.json_response({
            "ok": True,
            "filename": os.path.basename(saved),
            "cached": False,
        })

    logger.info(
        "Oz ShotSplitter: registered /oz_shotsplitter/{active_clients,download_url}"
    )


register_routes()