"""Small ComfyUI REST client used by AutoShortsBot clip jobs."""

from __future__ import annotations

import json
import time
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any


class UrllibComfyUITransport:
    """HTTP transport for local/Tailscale ComfyUI instances."""

    def upload_image(self, base_url: str, path: Path, *, remote_name: str, overwrite: bool = True) -> dict[str, Any]:
        boundary = "----AutoShortsBotBoundary"
        body: list[bytes] = []

        def add(part: str | bytes) -> None:
            body.append(part.encode() if isinstance(part, str) else part)

        add(
            f'--{boundary}\r\nContent-Disposition: form-data; name="image"; filename="{remote_name}"\r\n'
            "Content-Type: image/png\r\n\r\n"
        )
        add(path.read_bytes())
        add("\r\n")
        add(f'--{boundary}\r\nContent-Disposition: form-data; name="type"\r\n\r\ninput\r\n')
        add(f'--{boundary}\r\nContent-Disposition: form-data; name="overwrite"\r\n\r\n{str(overwrite).lower()}\r\n')
        add(f"--{boundary}--\r\n")
        request = urllib.request.Request(
            f"{base_url}/upload/image",
            data=b"".join(body),
            headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
        )
        return json.loads(urllib.request.urlopen(request, timeout=120).read().decode())

    def post_json(self, base_url: str, endpoint: str, payload: dict[str, Any]) -> dict[str, Any]:
        request = urllib.request.Request(
            f"{base_url}{endpoint}",
            data=json.dumps(payload).encode(),
            headers={"Content-Type": "application/json"},
        )
        raw = urllib.request.urlopen(request, timeout=120).read().decode()
        return json.loads(raw) if raw else {}

    def wait_for_completion(
        self,
        base_url: str,
        prompt_id: str,
        *,
        poll_seconds: float,
        timeout_seconds: int,
    ) -> dict[str, Any]:
        deadline = time.monotonic() + timeout_seconds
        while time.monotonic() < deadline:
            queue = json.loads(urllib.request.urlopen(f"{base_url}/queue", timeout=30).read().decode())
            running = queue.get("queue_running", [])
            pending = queue.get("queue_pending", [])
            if not running and not pending:
                history = json.loads(urllib.request.urlopen(f"{base_url}/history/{prompt_id}", timeout=60).read().decode())
                return history.get(prompt_id, {})
            time.sleep(poll_seconds)
        raise TimeoutError(f"ComfyUI prompt did not finish within {timeout_seconds}s: {prompt_id}")

    def download_view(self, base_url: str, output: dict[str, Any], destination: Path) -> Path:
        params = urllib.parse.urlencode(
            {
                "filename": output["filename"],
                "subfolder": output.get("subfolder", ""),
                "type": output.get("type", "output"),
            }
        )
        destination.parent.mkdir(parents=True, exist_ok=True)
        destination.write_bytes(urllib.request.urlopen(f"{base_url}/view?{params}", timeout=300).read())
        return destination


class ComfyUIClient:
    """Minimal client for uploading keyframes, running workflows, and downloading MP4 outputs."""

    def __init__(self, base_url: str, *, transport: Any | None = None):
        self.base_url = base_url.rstrip("/")
        self.transport = transport or UrllibComfyUITransport()

    def run_i2v_workflow(
        self,
        *,
        workflow: dict[str, Any],
        keyframe_path: Path,
        remote_image_name: str,
        output_path: Path,
        poll_seconds: float = 5,
        timeout_seconds: int = 900,
    ) -> dict[str, Any]:
        upload = self.transport.upload_image(self.base_url, keyframe_path, remote_name=remote_image_name, overwrite=True)
        prompt_response = self.transport.post_json(
            self.base_url,
            "/prompt",
            {"prompt": workflow, "client_id": "autoshorts-i2v"},
        )
        prompt_id = prompt_response["prompt_id"]
        history = self.transport.wait_for_completion(
            self.base_url,
            prompt_id,
            poll_seconds=poll_seconds,
            timeout_seconds=timeout_seconds,
        )
        status = history.get("status", {})
        if status.get("status_str") != "success":
            return {"status": "failed", "prompt_id": prompt_id, "history": history, "upload": upload}
        output = _first_video_output(history)
        self.transport.download_view(self.base_url, output, output_path)
        return {
            "status": "success",
            "prompt_id": prompt_id,
            "output_path": str(output_path),
            "remote_output": output,
            "upload": upload,
        }


def _first_video_output(history: dict[str, Any]) -> dict[str, Any]:
    for node_output in history.get("outputs", {}).values():
        for item in node_output.get("images", []):
            if str(item.get("filename", "")).lower().endswith(".mp4"):
                return item
    raise ValueError("No MP4 output found in ComfyUI history")