#!/usr/bin/env python3
"""
ComfyDock local agent — Phase 4

Tiny HTTP server that sits between the ComfyDock web app and your local
ComfyUI install. Phase 4 adds an outputs library (with workflow sidecars
for one-click re-run), resumable model downloads, SHA256 verification,
and an in-app Civitai search proxy.

Usage:
    # Windows
    py -m pip install fastapi uvicorn httpx python-multipart huggingface_hub
    py comfydock_agent.py --comfy-path "C:\\Users\\you\\Documents\\ComfyUI"

    # macOS / Linux
    python3 -m pip install --user fastapi uvicorn httpx python-multipart huggingface_hub
    python3 comfydock_agent.py --comfy-path ~/ComfyUI

On first run it prints an access token. Paste it into ComfyDock → Settings.
"""
from __future__ import annotations

import argparse
import asyncio
import atexit
import hashlib
import json
import mimetypes
import os
import secrets
import subprocess
import sys

# Force UTF-8 stdout/stderr on Windows so unicode chars (arrows, box drawing,
# em-dashes) in our banners don't crash with UnicodeEncodeError under cp1252.
try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")
except Exception:
    pass
import webbrowser
import uuid
from pathlib import Path
from typing import Any, Optional

try:
    import httpx
    import uvicorn
    from fastapi import Depends, FastAPI, File, HTTPException, Request, UploadFile
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.responses import FileResponse
except ImportError:
    sys.exit(
        "Missing deps. Run:\n"
        "  pip install fastapi uvicorn httpx python-multipart huggingface_hub"
    )

VERSION = "0.7.2"
CONFIG_DIR = Path.home() / ".comfydock"
CONFIG_FILE = CONFIG_DIR / "config.json"
HASHES_FILE = CONFIG_DIR / "hashes.json"
COMFY_URL = "http://127.0.0.1:8188"

MODEL_SUBDIRS = [
    "checkpoints", "loras", "vae", "controlnet",
    "embeddings", "upscale_models", "clip", "clip_vision", "unet",
]


def _read_config() -> dict[str, Any]:
    CONFIG_DIR.mkdir(exist_ok=True)
    if CONFIG_FILE.exists():
        try:
            return json.loads(CONFIG_FILE.read_text())
        except Exception:
            return {}
    return {}


def _write_config(cfg: dict[str, Any]) -> None:
    CONFIG_DIR.mkdir(exist_ok=True)
    CONFIG_FILE.write_text(json.dumps(cfg, indent=2))


def load_or_create_token() -> str:
    cfg = _read_config()
    token = cfg.get("token")
    if not token:
        token = secrets.token_urlsafe(24)
        cfg["token"] = token
        _write_config(cfg)
    return token


def get_civitai_token() -> Optional[str]:
    return _read_config().get("civitai_token") or None


def _read_hashes() -> dict[str, str]:
    if HASHES_FILE.exists():
        try:
            return json.loads(HASHES_FILE.read_text())
        except Exception:
            return {}
    return {}


def _write_hashes(h: dict[str, str]) -> None:
    CONFIG_DIR.mkdir(exist_ok=True)
    HASHES_FILE.write_text(json.dumps(h, indent=2))


def _set_expected_hash(folder: str, filename: str, sha256: str) -> None:
    h = _read_hashes()
    h[f"{folder}/{filename}"] = sha256.lower()
    _write_hashes(h)


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--comfy-path", default=None,
                   help="Path to your ComfyUI install (contains models/, output/). "
                        "If omitted, the agent auto-discovers ComfyUI on this machine.")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=7820)
    p.add_argument("--allow-origin", default="*",
                   help="CORS origin. Use the deployed ComfyDock URL for safety.")
    p.add_argument("--cloud-url", default="https://xeukvglgxnbbttkvpnkg.supabase.co",
                   help="Lovable Cloud base URL. Leave default for hosted ComfyDock.")
    p.add_argument("--cloud-key", default="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InhldWt2Z2xneG5iYnR0a3ZwbmtnIiwicm9sZSI6ImFub24iLCJpYXQiOjE3NzkxMjU4MTUsImV4cCI6MjA5NDcwMTgxNX0.0CTga7sHXeApHQYgRw01xxFLZeyN7xWT8C5UoTMPGGM",
                   help="Public API key for Lovable Cloud (safe to ship).")
    p.add_argument("--no-cloud", action="store_true",
                   help="Disable remote access (loopback HTTP only).")
    p.add_argument("--device-label", default=None,
                   help="Friendly name shown in the Devices page (defaults to hostname).")
    p.add_argument("--no-comfy", action="store_true",
                   help="Don't launch ComfyUI — assume it's already running on :8188.")
    p.add_argument("--comfy-python", default=None,
                   help="Python executable to use for ComfyUI. Defaults to portable embedded Python when found, otherwise this Python.")
    p.add_argument("--no-open", action="store_true",
                   help="Don't open the ComfyDock web app in your browser on launch.")
    p.add_argument("--web-url",
                   default="https://id-preview--1b0df54d-2b45-47a4-98a8-4ed68ed857a3.lovable.app",
                   help="ComfyDock web app URL to open in the browser.")
    p.add_argument("--print-pair-code", action="store_true",
                   help="Print the current cloud pair code and exit.")
    return p.parse_args()


args = parse_args()
TOKEN = load_or_create_token()


def _looks_like_comfy(folder: Path) -> bool:
    """A folder counts as ComfyUI only with multiple ComfyUI-specific markers."""
    try:
        main_py = folder / "main.py"
        if not main_py.is_file():
            return False
    except OSError:
        return False
    score = 0
    # main.py content signature
    try:
        head = main_py.read_text(encoding="utf-8", errors="ignore")[:4096].lower()
        if "comfyui" in head and ("--listen" in head or "server" in head):
            score += 1
    except OSError:
        pass
    # Core ComfyUI files/folders
    for marker in ("nodes.py", "folder_paths.py", "execution.py"):
        try:
            if (folder / marker).is_file():
                score += 1
        except OSError:
            pass
    try:
        if (folder / "comfy").is_dir():
            score += 1
    except OSError:
        pass
    try:
        if (folder / "web" / "index.html").is_file():
            score += 1
    except OSError:
        pass
    return score >= 2


def _path_has_comfy(p: Path) -> Optional[Path]:
    """Return the folder that is a real ComfyUI install, or None."""
    try:
        if not p.exists():
            return None
    except OSError:
        return None
    candidates = [
        p,
        p / "ComfyUI",
        p / "ComfyUI_windows_portable" / "ComfyUI",
    ]
    for c in candidates:
        try:
            if _looks_like_comfy(c):
                return c.resolve()
        except OSError:
            continue
    return None


def _windows_drive_roots() -> list[Path]:
    roots: list[Path] = []
    if os.name != "nt":
        return roots
    import string
    for letter in string.ascii_uppercase:
        root = Path(f"{letter}:/")
        try:
            if root.exists():
                roots.append(root)
        except OSError:
            pass
    return roots


def _candidate_search_dirs() -> list[Path]:
    """Common parent folders to look inside for a ComfyUI install."""
    home = Path.home()
    bases: list[Path] = [
        home, home / "Documents", home / "Desktop", home / "Downloads",
        home / "AI", home / "AI Tools", home / "Apps",
    ]
    if os.name == "nt":
        for root in _windows_drive_roots():
            bases.extend([
                root, root / "AI", root / "AI Tools", root / "Apps",
                root / "Tools", root / "ComfyUI", root / "StabilityMatrix",
            ])
    else:
        for root in [Path("/opt"), Path("/srv"), Path("/mnt"), Path("/media")]:
            if root.exists():
                bases.append(root)
    seen: set[str] = set()
    out: list[Path] = []
    for b in bases:
        try:
            key = str(b.resolve())
        except OSError:
            continue
        if key in seen:
            continue
        seen.add(key)
        out.append(b)
    return out


def _scan_for_comfy(base: Path, max_depth: int = 2) -> Optional[Path]:
    """Bounded scan: check `base` and immediate subfolders for ComfyUI."""
    found = _path_has_comfy(base)
    if found:
        return found
    if max_depth <= 0:
        return None
    try:
        children = list(base.iterdir())
    except (OSError, PermissionError):
        return None
    # Heuristic: prefer entries whose name hints at ComfyUI to keep the scan fast.
    hinted = [c for c in children if c.is_dir() and "comfy" in c.name.lower()]
    others = [c for c in children if c.is_dir() and c not in hinted]
    for child in hinted + others[:40]:  # cap breadth
        found = _scan_for_comfy(child, max_depth - 1)
        if found:
            return found
    return None


def resolve_comfy_path() -> Path:
    """Resolve the ComfyUI install path. Priority:
    1. --comfy-path CLI arg (used as-is, with portable/wrapper detection)
    2. saved path in ~/.comfydock/config.json
    3. auto-discovery across common locations and drives
    4. fallback to ~/ComfyUI so the rest of the agent still boots
    """
    tried: list[str] = []

    if args.comfy_path:
        raw = Path(args.comfy_path).expanduser().resolve()
        tried.append(str(raw))
        found = _path_has_comfy(raw)
        if found:
            return found
        print(f"  Comfy path: {raw} does not contain main.py — searching…", flush=True)

    cfg = _read_config()
    saved = cfg.get("comfy_path")
    if saved:
        sp = Path(saved).expanduser()
        tried.append(str(sp))
        found = _path_has_comfy(sp)
        if found:
            return found

    print("  Comfy path: searching for ComfyUI…", flush=True)
    for base in _candidate_search_dirs():
        found = _scan_for_comfy(base, max_depth=2)
        if found:
            print(f"  Comfy path: found at {found}", flush=True)
            return found

    fallback = (Path.home() / "ComfyUI").resolve()
    print("  Comfy path: NOT FOUND. Tried:", flush=True)
    for t in tried:
        print(f"              - {t}", flush=True)
    print("              Pass --comfy-path explicitly, e.g.:", flush=True)
    print('                --comfy-path "D:\\AI\\ComfyUI"', flush=True)
    print('                --comfy-path "E:\\ComfyUI_windows_portable"', flush=True)
    return fallback


COMFY_PATH = resolve_comfy_path()

# Persist the resolved path so the next run is instant.
try:
    if (COMFY_PATH / "main.py").is_file():
        _cfg = _read_config()
        if _cfg.get("comfy_path") != str(COMFY_PATH):
            _cfg["comfy_path"] = str(COMFY_PATH)
            _write_config(_cfg)
except Exception:
    pass

app = FastAPI(title="ComfyDock Agent", version=VERSION)
app.add_middleware(
    CORSMiddleware,
    allow_origins=[args.allow_origin] if args.allow_origin != "*" else ["*"],
    allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
    allow_headers=["*"],
)


def require_token(request: Request):
    # Token enforcement is soft in Phase 1 to make local testing painless.
    # Set REQUIRE_TOKEN=1 to enforce.
    if os.environ.get("REQUIRE_TOKEN") != "1":
        return
    auth = request.headers.get("authorization", "")
    if auth != f"Bearer {TOKEN}":
        raise HTTPException(401, "Invalid or missing token")


@app.get("/health")
async def health(_=Depends(require_token)):
    comfy_reachable = False
    comfy_version: Optional[str] = None
    try:
        async with httpx.AsyncClient(timeout=2.0) as client:
            r = await client.get(f"{COMFY_URL}/system_stats")
            if r.status_code == 200:
                comfy_reachable = True
                data = r.json()
                comfy_version = (
                    data.get("system", {}).get("comfyui_version")
                    or data.get("system", {}).get("python_version")
                )
    except Exception:
        pass
    return {
        "version": VERSION,
        "device_id": _read_config().get("cloud_device_id"),
        "comfy_path": str(COMFY_PATH),
        "comfy_reachable": comfy_reachable,
        "comfy_version": comfy_version,
    }


@app.get("/models")
async def list_models(_=Depends(require_token)):
    root = COMFY_PATH / "models"
    out = {}
    for sub in MODEL_SUBDIRS:
        d = root / sub
        items = []
        if d.is_dir():
            for f in sorted(d.iterdir()):
                if f.is_file() and not f.name.startswith("."):
                    items.append({"name": f.name, "size": f.stat().st_size})
        out[sub] = items
    return out


@app.get("/")
async def index():
    return {"name": "ComfyDock Agent", "version": VERSION}


# ---- Phase 2 endpoints --------------------------------------------------

@app.post("/upload")
async def upload(file: UploadFile = File(...), _=Depends(require_token)):
    input_dir = COMFY_PATH / "input"
    input_dir.mkdir(parents=True, exist_ok=True)
    data = await file.read()
    digest = hashlib.sha1(data).hexdigest()[:12]
    suffix = Path(file.filename or "upload.png").suffix or ".png"
    target = input_dir / f"comfydock_{digest}{suffix}"
    target.write_bytes(data)
    return {"filename": target.name}


@app.post("/workflows/run")
async def run_workflow(req: Request, _=Depends(require_token)):
    body = await req.json()
    prompt = body.get("prompt")
    workflow_id = body.get("workflow_id")
    inputs = body.get("inputs") if isinstance(body.get("inputs"), dict) else None
    if not isinstance(prompt, dict):
        raise HTTPException(400, "Missing prompt graph")
    client_id = secrets.token_hex(8)
    payload = {"prompt": prompt, "client_id": client_id}
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            r = await client.post(f"{COMFY_URL}/prompt", json=payload)
            if r.status_code != 200:
                raise HTTPException(502, f"ComfyUI rejected prompt: {r.text}")
            data = r.json()
            prompt_id = data.get("prompt_id")
            if prompt_id and (workflow_id or inputs):
                PROMPT_META[prompt_id] = {
                    "workflow_id": workflow_id,
                    "inputs": inputs or {},
                    "started_at": int(asyncio.get_event_loop().time() * 1000),
                }
            if prompt_id:
                asyncio.create_task(_track_run(prompt_id))
            return {"prompt_id": prompt_id}
    except httpx.HTTPError as e:
        raise HTTPException(502, f"Cannot reach ComfyUI: {e}")


@app.get("/workflows/run/{prompt_id}/status")
async def run_status(prompt_id: str, _=Depends(require_token)):
    """Polled by the web app. Returns queued/running/done/error + outputs."""
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            hist = await client.get(f"{COMFY_URL}/history/{prompt_id}")
            if hist.status_code == 200:
                hdata = hist.json()
                entry = hdata.get(prompt_id)
                if entry:
                    status_info = entry.get("status", {})
                    if status_info.get("status_str") == "error":
                        return {"status": "error", "error": _extract_error(status_info)}
                    outputs = []
                    for node_id, node_out in (entry.get("outputs") or {}).items():
                        for img in node_out.get("images", []) or []:
                            outputs.append({
                                "filename": img.get("filename"),
                                "subfolder": img.get("subfolder", ""),
                                "type": img.get("type", "output"),
                            })
                    if outputs:
                        _write_sidecars(prompt_id, outputs)
                        return {"status": "done", "outputs": outputs}

            # Still running: pull queue + progress
            q = await client.get(f"{COMFY_URL}/queue")
            queue = q.json() if q.status_code == 200 else {}
            running = queue.get("queue_running", []) or []
            pending = queue.get("queue_pending", []) or []
            if any(_match_prompt(item, prompt_id) for item in running):
                return {"status": "running", "step": 0, "total": 0}
            if any(_match_prompt(item, prompt_id) for item in pending):
                return {"status": "queued", "step": 0, "total": 0}
            # Not in queue, not in history yet — treat as queued briefly
            return {"status": "queued", "step": 0, "total": 0}
    except httpx.HTTPError as e:
        return {"status": "error", "error": str(e)}


def _match_prompt(item, prompt_id: str) -> bool:
    # Queue entries are tuples [number, prompt_id, prompt, ...]
    try:
        return item[1] == prompt_id
    except Exception:
        return False


def _extract_error(status_info) -> str:
    msgs = status_info.get("messages") or []
    for m in msgs:
        if isinstance(m, list) and len(m) >= 2 and m[0] == "execution_error":
            data = m[1] or {}
            return data.get("exception_message") or "Execution error"
    return "Execution failed"


@app.post("/workflows/run/{prompt_id}/cancel")
async def cancel(prompt_id: str, _=Depends(require_token)):
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            await client.post(f"{COMFY_URL}/interrupt")
    except httpx.HTTPError:
        pass
    return {"ok": True, "prompt_id": prompt_id}


@app.get("/outputs")
async def outputs(filename: str, subfolder: str = "", type: str = "output",
                  _=Depends(require_token)):
    # Safety: confine to ComfyUI's output/temp/input dirs and forbid traversal.
    base_map = {
        "output": COMFY_PATH / "output",
        "temp": COMFY_PATH / "temp",
        "input": COMFY_PATH / "input",
    }
    base = base_map.get(type, COMFY_PATH / "output")
    candidate = (base / subfolder / filename).resolve()
    try:
        candidate.relative_to(base.resolve())
    except ValueError:
        raise HTTPException(400, "Invalid path")
    if not candidate.is_file():
        raise HTTPException(404, "Not found")
    media_type = mimetypes.guess_type(str(candidate))[0] or "application/octet-stream"
    return FileResponse(candidate, media_type=media_type, headers={"Cache-Control": "no-store"})


# ---- Phase 3: model installer -------------------------------------------

INSTALLS: dict[str, dict[str, Any]] = {}
PROMPT_META: dict[str, dict[str, Any]] = {}


def _write_sidecars(prompt_id: str, outputs: list[dict]) -> None:
    meta = PROMPT_META.pop(prompt_id, None)
    if not meta:
        return
    out_dir = COMFY_PATH / "output"
    cfg = _read_config()
    device_id = cfg.get("cloud_device_id")
    payload = {
        "schema": "comfydock/1",
        "workflow_id": meta.get("workflow_id"),
        "inputs": meta.get("inputs") or {},
        "started_at": meta.get("started_at"),
        "finished_at": int(asyncio.get_event_loop().time() * 1000),
        "agent_version": VERSION,
        "prompt_id": prompt_id,
        "device_id": device_id,
    }
    for o in outputs:
        fn = o.get("filename")
        sub = o.get("subfolder") or ""
        if not fn:
            continue
        target = (out_dir / sub / fn).with_suffix(".json")
        try:
            target.parent.mkdir(parents=True, exist_ok=True)
            target.write_text(json.dumps(payload, indent=2))
        except Exception:
            pass


def _safe_folder(folder: str) -> Path:
    if folder not in MODEL_SUBDIRS:
        raise HTTPException(400, f"Unknown folder: {folder}")
    p = COMFY_PATH / "models" / folder
    p.mkdir(parents=True, exist_ok=True)
    return p


def _safe_filename(name: str) -> str:
    base = Path(name).name
    if not base or base.startswith("."):
        raise HTTPException(400, "Invalid filename")
    return base


async def _stream_download(
    install_id: str, url: str, dest: Path, extra_headers: dict[str, str] | None = None
):
    state = INSTALLS[install_id]
    part = dest.with_suffix(dest.suffix + ".part")
    resume_from = 0
    if part.exists():
        try:
            resume_from = part.stat().st_size
        except Exception:
            resume_from = 0
    req_headers = dict(extra_headers or {})
    if resume_from > 0:
        req_headers["Range"] = f"bytes={resume_from}-"
    try:
        async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client:
            async with client.stream("GET", url, headers=req_headers) as r:
                if r.status_code >= 400:
                    # 416 = invalid range (already complete); fall through to rename
                    if r.status_code == 416 and resume_from > 0:
                        part.replace(dest)
                        state["status"] = "done"
                        state["filename"] = dest.name
                        return
                    raise HTTPException(502, f"Upstream {r.status_code}")
                resumed = r.status_code == 206 and resume_from > 0
                cl = int(r.headers.get("content-length", "0") or 0)
                if resumed:
                    state["total_bytes"] = resume_from + cl
                    state["resumed"] = True
                elif cl:
                    state["total_bytes"] = cl
                downloaded = resume_from if resumed else 0
                state["downloaded_bytes"] = downloaded
                mode = "ab" if resumed else "wb"
                with part.open(mode) as f:
                    async for chunk in r.aiter_bytes(chunk_size=1 << 20):
                        if state.get("cancelled"):
                            f.close()
                            try:
                                part.unlink()
                            except Exception:
                                pass
                            state["status"] = "cancelled"
                            return
                        f.write(chunk)
                        downloaded += len(chunk)
                        state["downloaded_bytes"] = downloaded
        part.replace(dest)
        state["status"] = "done"
        state["filename"] = dest.name
        # Record expected hash if known
        folder = state.get("folder")
        expected = state.get("expected_sha256")
        if folder and expected:
            try:
                _set_expected_hash(folder, dest.name, expected)
            except Exception:
                pass
    except HTTPException as e:
        state["status"] = "error"
        state["error"] = e.detail
    except Exception as e:  # noqa: BLE001
        state["status"] = "error"
        state["error"] = str(e)


async def _civitai_resolve(payload: dict) -> tuple[str, str, dict[str, str]]:
    token = get_civitai_token()
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    version_id = payload.get("version_id")
    url = (payload.get("url") or "").strip()
    if not version_id and url and "civitai.com" in url:
        from urllib.parse import urlparse, parse_qs
        q = parse_qs(urlparse(url).query)
        if "modelVersionId" in q:
            version_id = int(q["modelVersionId"][0])
    if version_id:
        api = f"https://civitai.com/api/v1/model-versions/{version_id}"
        async with httpx.AsyncClient(timeout=15.0, follow_redirects=True) as client:
            r = await client.get(api, headers=headers)
            if r.status_code != 200:
                raise HTTPException(502, f"Civitai API error {r.status_code}")
            data = r.json()
            files = data.get("files") or []
            if not files:
                raise HTTPException(404, "Civitai version has no files")
            target = files[0]
            sha = ((target.get("hashes") or {}).get("SHA256") or "").lower() or None
            return (
                target["downloadUrl"],
                payload.get("filename") or target.get("name") or f"civitai_{version_id}.safetensors",
                {**headers, "_sha256": sha} if sha else headers,
            )
    if url:
        return url, payload.get("filename") or "civitai_download.safetensors", headers
    raise HTTPException(400, "Missing Civitai url or version_id")


async def _hf_download(install_id: str, repo_id: str, filename: str, dest_dir: Path):
    state = INSTALLS[install_id]
    try:
        from huggingface_hub import hf_hub_url, HfFileSystem  # type: ignore
    except ImportError:
        state["status"] = "error"
        state["error"] = "huggingface_hub not installed. pip install huggingface_hub"
        return
    url = hf_hub_url(repo_id=repo_id, filename=filename)
    dest = dest_dir / _safe_filename(filename)
    try:
        info = HfFileSystem().info(f"{repo_id}/{filename}")
        if isinstance(info, dict) and info.get("size"):
            state["total_bytes"] = int(info["size"])
    except Exception:
        pass
    await _stream_download(install_id, url, dest)


async def _run_install(install_id: str, payload: dict):
    state = INSTALLS[install_id]
    try:
        folder = _safe_folder(payload["target_folder"])
        state["folder"] = payload["target_folder"]
        source = payload.get("source")
        if source == "huggingface":
            await _hf_download(install_id, payload["repo_id"], payload["filename"], folder)
        elif source == "civitai":
            url, filename, headers = await _civitai_resolve(payload)
            sha = headers.pop("_sha256", None) if isinstance(headers, dict) else None
            if sha:
                state["expected_sha256"] = sha
            dest = folder / _safe_filename(filename)
            await _stream_download(install_id, url, dest, extra_headers=headers)
        elif source == "url":
            dest = folder / _safe_filename(payload["filename"])
            await _stream_download(install_id, payload["url"], dest)
        else:
            state["status"] = "error"
            state["error"] = f"Unknown source: {source}"
    except HTTPException as e:
        state["status"] = "error"
        state["error"] = e.detail
    except Exception as e:  # noqa: BLE001
        state["status"] = "error"
        state["error"] = str(e)


@app.post("/models/install")
async def models_install(req: Request, _=Depends(require_token)):
    payload = await req.json()
    install_id = uuid.uuid4().hex
    INSTALLS[install_id] = {
        "status": "downloading",
        "downloaded_bytes": 0,
        "total_bytes": 0,
        "cancelled": False,
    }
    asyncio.create_task(_run_install(install_id, payload))
    return {"install_id": install_id}


@app.get("/models/install/{install_id}/status")
async def models_install_status(install_id: str, _=Depends(require_token)):
    s = INSTALLS.get(install_id)
    if not s:
        raise HTTPException(404, "Unknown install id")
    return {
        "status": s.get("status", "downloading"),
        "downloaded_bytes": s.get("downloaded_bytes", 0),
        "total_bytes": s.get("total_bytes", 0),
        "filename": s.get("filename"),
        "error": s.get("error"),
        "resumed": bool(s.get("resumed")),
    }


@app.post("/models/install/{install_id}/cancel")
async def models_install_cancel(install_id: str, _=Depends(require_token)):
    s = INSTALLS.get(install_id)
    if not s:
        raise HTTPException(404, "Unknown install id")
    s["cancelled"] = True
    return {"ok": True}


@app.delete("/models/{folder}/{filename}")
async def models_delete(folder: str, filename: str, _=Depends(require_token)):
    dest_dir = _safe_folder(folder)
    target = (dest_dir / _safe_filename(filename)).resolve()
    try:
        target.relative_to(dest_dir.resolve())
    except ValueError:
        raise HTTPException(400, "Invalid path")
    if not target.is_file():
        raise HTTPException(404, "Not found")
    target.unlink()
    return {"ok": True}


@app.post("/config")
async def set_config(req: Request, _=Depends(require_token)):
    incoming = await req.json()
    cfg = _read_config()
    if "civitai_token" in incoming:
        cfg["civitai_token"] = (incoming.get("civitai_token") or "").strip() or None
    _write_config(cfg)
    return {"ok": True}


# ---- Custom node packs (auto-install from workflows) -------------------

NODE_INSTALLS: dict[str, dict[str, Any]] = {}


def _custom_nodes_dir() -> Path:
    p = COMFY_PATH / "custom_nodes"
    p.mkdir(parents=True, exist_ok=True)
    return p


def _safe_node_dir(name: str) -> str:
    base = Path(name).name
    if not base or base.startswith(".") or "/" in name or "\\" in name:
        raise HTTPException(400, "Invalid node dir")
    return base


@app.get("/nodes")
async def list_nodes(_=Depends(require_token)):
    root = _custom_nodes_dir()
    nodes: list[dict] = []
    for child in sorted(root.iterdir()):
        if not child.is_dir():
            continue
        if child.name.startswith(".") or child.name == "__pycache__":
            continue
        nodes.append({"dir": child.name})
    return {"nodes": nodes}


async def _run_node_install(install_id: str, payload: dict):
    state = NODE_INSTALLS[install_id]
    try:
        git_url = (payload.get("git_url") or "").strip()
        if not git_url or not (git_url.startswith("https://") or git_url.startswith("http://") or git_url.startswith("git@")):
            state["status"] = "error"
            state["error"] = "Invalid git_url"
            return
        target = _safe_node_dir(payload.get("dir") or "")
        ref = (payload.get("ref") or "").strip() or None
        install_reqs = payload.get("install_requirements", True)
        dest = _custom_nodes_dir() / target
        if dest.exists():
            state["status"] = "done"
            state["message"] = "Already installed"
            return
        state["status"] = "cloning"
        state["message"] = f"git clone {git_url}"
        proc = await asyncio.create_subprocess_exec(
            "git", "clone", "--depth", "1", git_url, str(dest),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        _, err = await proc.communicate()
        if proc.returncode != 0:
            state["status"] = "error"
            state["error"] = (err.decode(errors="ignore") or "git clone failed").strip()
            return
        if ref:
            state["message"] = f"checkout {ref}"
            proc = await asyncio.create_subprocess_exec(
                "git", "-C", str(dest), "fetch", "--depth", "1", "origin", ref,
                stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE,
            )
            await proc.communicate()
            proc = await asyncio.create_subprocess_exec(
                "git", "-C", str(dest), "checkout", ref,
                stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE,
            )
            await proc.communicate()
        if install_reqs:
            reqs = dest / "requirements.txt"
            if reqs.is_file():
                state["status"] = "installing_deps"
                state["message"] = "pip install -r requirements.txt"
                proc = await asyncio.create_subprocess_exec(
                    sys.executable, "-m", "pip", "install", "-r", str(reqs),
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                _, err = await proc.communicate()
                if proc.returncode != 0:
                    state["status"] = "error"
                    state["error"] = (err.decode(errors="ignore") or "pip install failed").strip()[-500:]
                    return
        state["status"] = "done"
        state["message"] = "Installed. Restart ComfyUI to load the new nodes."
    except Exception as e:  # noqa: BLE001
        state["status"] = "error"
        state["error"] = str(e)


@app.post("/nodes/install")
async def nodes_install(req: Request, _=Depends(require_token)):
    payload = await req.json()
    install_id = uuid.uuid4().hex
    NODE_INSTALLS[install_id] = {"status": "cloning", "message": "starting…"}
    asyncio.create_task(_run_node_install(install_id, payload))
    return {"install_id": install_id}


@app.get("/nodes/install/{install_id}/status")
async def nodes_install_status(install_id: str, _=Depends(require_token)):
    s = NODE_INSTALLS.get(install_id)
    if not s:
        raise HTTPException(404, "Unknown install id")
    return {
        "status": s.get("status", "cloning"),
        "message": s.get("message"),
        "error": s.get("error"),
    }


@app.delete("/nodes/{dir_name}")
async def nodes_remove(dir_name: str, _=Depends(require_token)):
    """Safely delete a custom node pack folder under ComfyUI/custom_nodes/."""
    import shutil
    base = _custom_nodes_dir().resolve()
    safe = _safe_node_dir(dir_name)
    target = (base / safe).resolve()
    # Path containment guard: target must be a direct child of base.
    try:
        if target.parent != base:
            raise HTTPException(400, "Invalid path")
    except Exception:
        raise HTTPException(400, "Invalid path")
    if not target.exists() or not target.is_dir():
        raise HTTPException(404, "Not installed")
    # Refuse to follow symlinks that escape custom_nodes/
    try:
        if target.is_symlink():
            raise HTTPException(400, "Refusing to remove symlink")
    except OSError:
        pass
    try:
        shutil.rmtree(target)
    except Exception as e:  # noqa: BLE001
        raise HTTPException(500, f"Remove failed: {e}")
    return {"ok": True, "dir": safe}


@app.post("/nodes/restart_comfy")
async def nodes_restart_comfy(_=Depends(require_token)):
    """Best-effort restart of the ComfyUI process this agent launched.

    If the agent didn't launch ComfyUI itself (--no-comfy) we don't try to
    kill an unrelated process; we return 501 and the UI shows manual instructions.
    """
    proc = globals().get("COMFY_PROC")
    launcher = globals().get("_relaunch_comfy")
    if proc is None or launcher is None:
        raise HTTPException(501, "Agent doesn't manage the ComfyUI process")
    try:
        try:
            proc.terminate()
        except Exception:
            pass
        # Give it a moment to shut down
        for _ in range(20):
            if proc.poll() is not None:
                break
            await asyncio.sleep(0.25)
        if proc.poll() is None:
            try:
                proc.kill()
            except Exception:
                pass
        new_proc = launcher()
        globals()["COMFY_PROC"] = new_proc
        return {"ok": True}
    except Exception as e:  # noqa: BLE001
        raise HTTPException(500, f"Restart failed: {e}")


# ---- Phase 4: outputs library, verify, civitai search ------------------

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".gif"}


@app.get("/outputs/index")
async def outputs_index(_=Depends(require_token)):
    out_dir = COMFY_PATH / "output"
    items: list[dict] = []
    if out_dir.is_dir():
        for p in out_dir.rglob("*"):
            if not p.is_file():
                continue
            if p.suffix.lower() not in IMAGE_EXTS:
                continue
            try:
                st = p.stat()
            except Exception:
                continue
            sub = str(p.parent.relative_to(out_dir)).replace("\\", "/")
            if sub == ".":
                sub = ""
            sidecar = p.with_suffix(".json").is_file()
            items.append({
                "filename": p.name,
                "subfolder": sub,
                "size": st.st_size,
                "mtime": int(st.st_mtime * 1000),
                "has_sidecar": sidecar,
            })
    items.sort(key=lambda x: x["mtime"], reverse=True)
    return {"items": items[:500]}


@app.get("/outputs/sidecar")
async def outputs_sidecar(filename: str, subfolder: str = "",
                          _=Depends(require_token)):
    base = COMFY_PATH / "output"
    candidate = (base / subfolder / filename).resolve()
    try:
        candidate.relative_to(base.resolve())
    except ValueError:
        raise HTTPException(400, "Invalid path")
    sidecar = candidate.with_suffix(".json")
    if not sidecar.is_file():
        raise HTTPException(404, "No sidecar")
    try:
        return json.loads(sidecar.read_text())
    except Exception as e:  # noqa: BLE001
        raise HTTPException(500, f"Bad sidecar: {e}")


@app.get("/models/{folder}/{filename}/verify")
async def models_verify(folder: str, filename: str, _=Depends(require_token)):
    dest_dir = _safe_folder(folder)
    target = (dest_dir / _safe_filename(filename)).resolve()
    try:
        target.relative_to(dest_dir.resolve())
    except ValueError:
        raise HTTPException(400, "Invalid path")
    if not target.is_file():
        raise HTTPException(404, "Not found")
    expected = _read_hashes().get(f"{folder}/{filename}")
    if not expected:
        return {"status": "unknown"}
    h = hashlib.sha256()
    loop = asyncio.get_event_loop()

    def _hash() -> str:
        with target.open("rb") as f:
            for chunk in iter(lambda: f.read(1 << 20), b""):
                h.update(chunk)
        return h.hexdigest()

    actual = await loop.run_in_executor(None, _hash)
    match = actual.lower() == expected.lower()
    return {
        "status": "match" if match else "mismatch",
        "expected": expected,
        "actual": actual,
    }


@app.get("/civitai/search")
async def civitai_search(q: str = "", type: str = "", limit: int = 20,
                         nsfw: bool = False, _=Depends(require_token)):
    if not q.strip():
        return {"items": []}
    token = get_civitai_token()
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    params: dict[str, Any] = {
        "query": q,
        "limit": max(1, min(int(limit), 50)),
        "nsfw": "true" if nsfw else "false",
    }
    if type:
        params["types"] = type
    try:
        async with httpx.AsyncClient(timeout=15.0, follow_redirects=True) as client:
            r = await client.get("https://civitai.com/api/v1/models",
                                 params=params, headers=headers)
            if r.status_code != 200:
                raise HTTPException(502, f"Civitai search failed: {r.status_code}")
            data = r.json()
    except httpx.HTTPError as e:
        raise HTTPException(502, f"Civitai search error: {e}")
    items = []
    for m in (data.get("items") or [])[:limit]:
        versions = []
        for v in (m.get("modelVersions") or [])[:3]:
            files = v.get("files") or []
            f0 = files[0] if files else {}
            versions.append({
                "id": v.get("id"),
                "name": v.get("name"),
                "size_kb": f0.get("sizeKB") or 0,
                "filename": f0.get("name"),
                "download_url": f0.get("downloadUrl") or v.get("downloadUrl"),
            })
        imgs = m.get("modelVersions", [{}])[0].get("images") or []
        thumb = imgs[0].get("url") if imgs else None
        items.append({
            "id": m.get("id"),
            "name": m.get("name"),
            "type": m.get("type"),
            "nsfw": m.get("nsfw", False),
            "creator": (m.get("creator") or {}).get("username"),
            "thumbnail": thumb,
            "versions": versions,
        })
    return {"items": items}



# ---- Phase 5: Cloud relay -------------------------------------------------

import base64
import platform
import threading
import time
import urllib.request
import urllib.error


def _rpc(url: str, key: str, fn: str, payload: dict) -> Any:
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        f"{url}/rest/v1/rpc/{fn}",
        data=body,
        method="POST",
        headers={
            "Content-Type": "application/json",
            "apikey": key,
            "Authorization": f"Bearer {key}",
            "Accept": "application/json",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=20) as r:
            data = r.read()
            if not data:
                return None
            return json.loads(data)
    except urllib.error.HTTPError as e:
        msg = e.read().decode("utf-8", "replace")
        raise RuntimeError(f"RPC {fn} failed: {e.code} {msg}")


def _register_device(url: str, key: str, label: str) -> dict:
    rows = _rpc(url, key, "register_device", {"p_label": label})
    if isinstance(rows, list) and rows:
        return rows[0]
    if isinstance(rows, dict):
        return rows
    raise RuntimeError("register_device returned empty")


def _print_pair_banner(code: str, label: str):
    box = [
        "",
        "",
        "  >>>>>>>>>>>>>>>>>>>  PAIRING CODE  <<<<<<<<<<<<<<<<<<<",
        "  ┌──────────────────────────────────────────────────┐",
        "  │  Pair this PC with ComfyDock                     │",
        "  │                                                  │",
        "  │  1. Open the ComfyDock web app and sign in       │",
        "  │  2. Go to Devices → Pair a new device            │",
        "  │  3. Enter this code:                             │",
        "  │                                                  │",
        f"  │      {code:<42}  │",
        "  │                                                  │",
        f"  │  Device label: {label:<32}  │",
        "  └──────────────────────────────────────────────────┘",
        "  ^^^ Type the COMFY-WORD-WORD-#### code above into the web app.",
        "  (Not the 'Access token' line — that's local-only.)",
        "",
        "",
    ]
    print("\n".join(box), flush=True)


PAIR_CODE_FILE = CONFIG_DIR / "pair-code.txt"
_LAST_PAIR_CODE: Optional[str] = None
_LAST_PAIR_LABEL: str = ""


def _write_pair_code_file(code: str, label: str) -> None:
    try:
        CONFIG_DIR.mkdir(exist_ok=True)
        PAIR_CODE_FILE.write_text(
            f"{code}\n\nDevice label: {label}\n"
            f"Enter this code in the ComfyDock web app under Devices -> Pair a new device.\n",
            encoding="utf-8",
        )
    except OSError:
        pass


def _clear_pair_code_file() -> None:
    try:
        if PAIR_CODE_FILE.exists():
            PAIR_CODE_FILE.unlink()
    except OSError:
        pass


def _dispatch_local(method: str, path: str, headers: dict, body_b64: Optional[str]) -> tuple[int, dict, Optional[str]]:
    """Forward a relayed request to the local FastAPI app over loopback HTTP."""
    target = f"http://127.0.0.1:{args.port}{path}"
    data = base64.b64decode(body_b64) if body_b64 else None
    fwd_headers = {k: v for k, v in (headers or {}).items()
                   if k.lower() not in ("host", "content-length", "connection")}
    # Relayed requests are already authenticated by the cloud (agent_token
    # via RPC). Inject the local access token so the loopback FastAPI app —
    # which requires Bearer auth on every endpoint — accepts the forwarded
    # request. Without this, cloud-routed calls to /outputs, /health, etc.
    # fail with 401 and thumbnails never render.
    fwd_headers["Authorization"] = f"Bearer {TOKEN}"
    req = urllib.request.Request(target, data=data, method=method, headers=fwd_headers)
    try:
        with urllib.request.urlopen(req, timeout=300) as r:
            payload = r.read()
            status = r.getcode()
            resp_headers = {
                k.lower(): v for k, v in r.getheaders()
                if k.lower() not in ("content-length", "transfer-encoding", "connection")
            }
            if path.startswith("/outputs?") and "content-type" not in resp_headers:
                resp_headers["content-type"] = mimetypes.guess_type(path.split("filename=", 1)[-1].split("&", 1)[0])[0] or "application/octet-stream"
    except urllib.error.HTTPError as e:
        payload = e.read()
        status = e.code
        resp_headers = {}
    except Exception as e:
        return 502, {"content-type": "application/json"}, base64.b64encode(
            json.dumps({"error": str(e)}).encode("utf-8")
        ).decode("ascii")
    return status, resp_headers, base64.b64encode(payload).decode("ascii") if payload else None


def _relay_loop(cloud_url: str, cloud_key: str, token: str):
    backoff = 1.0
    while True:
        try:
            # Heartbeat every loop iteration (bundled with dequeue side-effect too)
            rows = _rpc(cloud_url, cloud_key, "agent_dequeue", {"p_token": token, "p_limit": 8})
            if not rows:
                time.sleep(0.3)
                backoff = 1.0
                continue
            for row in rows:
                rid = row.get("request_id")
                method = (row.get("method") or "GET").upper()
                path = row.get("path") or "/"
                headers = row.get("headers") or {}
                body = row.get("body")
                try:
                    status, resp_headers, resp_body = _dispatch_local(method, path, headers, body)
                except Exception as e:
                    status, resp_headers, resp_body = 500, {}, base64.b64encode(
                        json.dumps({"error": str(e)}).encode("utf-8")
                    ).decode("ascii")
                try:
                    _rpc(cloud_url, cloud_key, "agent_respond", {
                        "p_token": token,
                        "p_request_id": rid,
                        "p_status": status,
                        "p_headers": resp_headers,
                        "p_body": resp_body,
                    })
                except Exception as e:
                    print(f"[cloud] respond failed for {rid}: {e}", flush=True)
        except Exception as e:
            print(f"[cloud] relay error: {e} (retry in {backoff:.0f}s)", flush=True)
            time.sleep(backoff)
            backoff = min(backoff * 2, 30)


def _heartbeat_loop(cloud_url: str, cloud_key: str, token: str):
    while True:
        try:
            _rpc(cloud_url, cloud_key, "agent_heartbeat", {"p_token": token, "p_version": VERSION})
        except Exception:
            pass
        time.sleep(20)


# Cloud config captured at relay start so async tasks can report progress.
CLOUD: dict[str, Optional[str]] = {"url": None, "key": None, "token": None}


def _report_run(prompt_id: str, patch: dict) -> None:
    url, key, token = CLOUD.get("url"), CLOUD.get("key"), CLOUD.get("token")
    if not url or not key or not token:
        return
    try:
        _rpc(url, key, "agent_report_run", {
            "p_token": token,
            "p_prompt_id": prompt_id,
            "p_patch": patch,
        })
    except Exception as e:
        print(f"[cloud] report failed for {prompt_id}: {e}", flush=True)


# ---- Thumbnail upload to Lovable Cloud ----------------------------------

WEB_URL_FOR_THUMBS = args.web_url.rstrip("/")
_PIL_WARNED = False


def _resolve_output_path(o: dict) -> Optional[Path]:
    base_map = {
        "output": COMFY_PATH / "output",
        "temp": COMFY_PATH / "temp",
        "input": COMFY_PATH / "input",
    }
    base = base_map.get(o.get("type") or "output", COMFY_PATH / "output")
    sub = o.get("subfolder") or ""
    fn = o.get("filename")
    if not fn:
        return None
    try:
        candidate = (base / sub / fn).resolve()
        candidate.relative_to(base.resolve())
    except (ValueError, OSError):
        return None
    return candidate if candidate.is_file() else None


def _make_thumbnail_jpeg(src: Path, max_side: int = 512, quality: int = 75) -> Optional[bytes]:
    global _PIL_WARNED
    try:
        from PIL import Image  # type: ignore
    except ImportError:
        if not _PIL_WARNED:
            print("[cloud] Pillow not installed — skipping cloud previews. "
                  "Run: pip install Pillow", flush=True)
            _PIL_WARNED = True
        return None
    import io
    try:
        with Image.open(src) as im:
            im = im.convert("RGB")
            im.thumbnail((max_side, max_side))
            buf = io.BytesIO()
            im.save(buf, format="JPEG", quality=quality, optimize=True)
            return buf.getvalue()
    except Exception as e:  # noqa: BLE001
        print(f"[cloud] thumbnail encode failed for {src.name}: {e}", flush=True)
        return None


def _upload_thumbnail(prompt_id: str, index: int, jpeg_bytes: bytes) -> Optional[str]:
    token = CLOUD.get("token")
    key = CLOUD.get("key")
    if not token or not key:
        return None
    # 1) Get a signed upload URL from the web app
    try:
        req = urllib.request.Request(
            f"{WEB_URL_FOR_THUMBS}/api/public/agent-upload-url",
            data=json.dumps({
                "device_token": token,
                "prompt_id": prompt_id,
                "index": index,
            }).encode("utf-8"),
            method="POST",
            headers={"Content-Type": "application/json"},
        )
        with urllib.request.urlopen(req, timeout=15) as r:
            info = json.loads(r.read())
    except Exception as e:
        print(f"[cloud] upload-url failed: {e}", flush=True)
        return None
    signed_url = info.get("signed_url")
    path = info.get("path")
    if not signed_url or not path:
        return None
    # 2) PUT the JPEG bytes to the signed URL
    try:
        put = urllib.request.Request(
            signed_url,
            data=jpeg_bytes,
            method="PUT",
            headers={
                "Content-Type": "image/jpeg",
                "x-upsert": "true",
                "Authorization": f"Bearer {key}",
            },
        )
        with urllib.request.urlopen(put, timeout=60) as r:
            if r.getcode() >= 300:
                return None
    except Exception as e:
        print(f"[cloud] thumbnail PUT failed: {e}", flush=True)
        return None
    return path


def _attach_thumbnail(prompt_id: str, filename: str, path: str) -> None:
    url = CLOUD.get("url"); key = CLOUD.get("key"); token = CLOUD.get("token")
    if not url or not key or not token:
        return
    try:
        _rpc(url, key, "agent_attach_thumbnail", {
            "p_token": token,
            "p_prompt_id": prompt_id,
            "p_filename": filename,
            "p_thumbnail_path": path,
        })
    except Exception as e:
        print(f"[cloud] attach_thumbnail failed: {e}", flush=True)


def _upload_run_thumbnails(prompt_id: str, outputs: list[dict]) -> None:
    """Encode + upload up to the first 4 image outputs as JPEG thumbnails."""
    capped = (outputs or [])[:4]
    for i, o in enumerate(capped):
        fn = o.get("filename") or ""
        if not fn or fn.lower().endswith((".mp4", ".webm", ".mov", ".m4v")):
            continue
        src = _resolve_output_path(o)
        if not src:
            continue
        jpeg = _make_thumbnail_jpeg(src)
        if not jpeg:
            return  # Pillow missing or encode failed for all
        path = _upload_thumbnail(prompt_id, i, jpeg)
        if path:
            _attach_thumbnail(prompt_id, fn, path)


async def _track_run(prompt_id: str) -> None:
    """Poll local /workflows/run/:id/status and mirror to cloud runs row."""
    reported_started = False
    while True:
        try:
            data = await run_status(prompt_id)  # type: ignore[arg-type]
        except Exception:
            await asyncio.sleep(1.0)
            continue
        status = data.get("status")
        if status in ("running", "queued"):
            if not reported_started and status == "running":
                _report_run(prompt_id, {"status": "running",
                                        "step": int(data.get("step") or 0),
                                        "total": int(data.get("total") or 0)})
                reported_started = True
            elif reported_started:
                _report_run(prompt_id, {"status": "running",
                                        "step": int(data.get("step") or 0),
                                        "total": int(data.get("total") or 0)})
            await asyncio.sleep(1.0)
            continue
        if status == "done":
            outs = data.get("outputs") or []
            _report_run(prompt_id, {"status": "done", "outputs": outs})
            # Fire-and-forget thumbnail upload so we don't block the report.
            try:
                threading.Thread(
                    target=_upload_run_thumbnails,
                    args=(prompt_id, outs),
                    daemon=True,
                ).start()
            except Exception as e:
                print(f"[cloud] thumbnail thread failed: {e}", flush=True)
            return
        if status == "error":
            _report_run(prompt_id, {"status": "error", "error": data.get("error") or "error"})
            return
        await asyncio.sleep(1.0)


def start_cloud_relay():
    global _LAST_PAIR_CODE, _LAST_PAIR_LABEL
    cfg = _read_config()
    cloud_url = args.cloud_url.rstrip("/")
    cloud_key = args.cloud_key
    label = args.device_label or platform.node() or "My PC"
    _LAST_PAIR_LABEL = label

    device_id = cfg.get("cloud_device_id")
    token = cfg.get("cloud_agent_token")
    pairing_code = cfg.get("cloud_pairing_code")
    paired = cfg.get("cloud_paired", False)

    if not device_id or not token:
        try:
            row = _register_device(cloud_url, cloud_key, label)
            device_id = row["device_id"]
            token = row["agent_token"]
            pairing_code = row["pairing_code"]
            cfg.update({
                "cloud_device_id": device_id,
                "cloud_agent_token": token,
                "cloud_pairing_code": pairing_code,
                "cloud_paired": False,
            })
            _write_config(cfg)
        except Exception as e:
            print("", flush=True)
            print("  ##################################################", flush=True)
            print("  #  CLOUD PAIRING FAILED — no pair code generated #", flush=True)
            print("  ##################################################", flush=True)
            print(f"  Error: {e}", flush=True)
            print("  Check your internet connection / firewall and restart the agent.", flush=True)
            print("  Local mode still works (loopback only, no remote access).", flush=True)
            print("", flush=True)
            return

    if not paired and pairing_code:
        _LAST_PAIR_CODE = pairing_code
        _write_pair_code_file(pairing_code, label)
        _print_pair_banner(pairing_code, label)
        print(f"  (also saved to: {PAIR_CODE_FILE})", flush=True)
        print("  (this banner won't print again once the device is paired)\n", flush=True)
    elif paired:
        _clear_pair_code_file()

    print(f"  Cloud relay: enabled (device {device_id[:8]}…)", flush=True)
    CLOUD["url"] = cloud_url
    CLOUD["key"] = cloud_key
    CLOUD["token"] = token
    threading.Thread(target=_relay_loop, args=(cloud_url, cloud_key, token), daemon=True).start()
    threading.Thread(target=_heartbeat_loop, args=(cloud_url, cloud_key, token), daemon=True).start()


def _resolve_comfy_main() -> Optional[Path]:
    candidates = [
        COMFY_PATH / "main.py",
        COMFY_PATH / "ComfyUI" / "main.py",
        COMFY_PATH / "ComfyUI_windows_portable" / "ComfyUI" / "main.py",
    ]
    for candidate in candidates:
        if candidate.is_file():
            return candidate
    return None


def _resolve_comfy_python(main_py: Path) -> str:
    if args.comfy_python:
        return str(Path(args.comfy_python).expanduser().resolve())
    roots = [COMFY_PATH, main_py.parent.parent, main_py.parent]
    names = ["python_embeded/python.exe", "python_embedded/python.exe", ".venv/Scripts/python.exe", ".venv/bin/python"]
    for root in roots:
        for name in names:
            candidate = root / name
            if candidate.is_file():
                return str(candidate)
    return sys.executable


def launch_comfyui() -> Optional[subprocess.Popen]:
    """Spawn ComfyUI as a child process without opening the ComfyUI browser UI."""
    main_py = _resolve_comfy_main()
    if main_py is None:
        print(f"  ComfyUI:    SKIP (no main.py under {COMFY_PATH})", flush=True)
        print("              Pass --comfy-path pointing to ComfyUI or ComfyUI_windows_portable.", flush=True)
        return None
    python_exe = _resolve_comfy_python(main_py)
    cmd = [python_exe, str(main_py),
           "--listen", "127.0.0.1",
           "--port", "8188"]
    try:
        proc = subprocess.Popen(cmd, cwd=str(main_py.parent))
    except Exception as e:
        print(f"  ComfyUI:    failed to start ({e})", flush=True)
        return None

    def _shutdown():
        try:
            if proc.poll() is None:
                proc.terminate()
                try:
                    proc.wait(timeout=5)
                except Exception:
                    proc.kill()
        except Exception:
            pass
    atexit.register(_shutdown)
    print(f"  ComfyUI:    starting backend (process id {proc.pid})", flush=True)
    print(f"              {python_exe} {main_py}", flush=True)
    print("              waiting for /system_stats…", flush=True)
    # Poll readiness up to ~30s.
    import urllib.request as _u
    import urllib.error as _ue
    import time as _t
    deadline = _t.time() + 30.0
    ready = False
    while _t.time() < deadline:
        if proc.poll() is not None:
            print(f"  ComfyUI:    exited early (code {proc.returncode})", flush=True)
            return None
        try:
            with _u.urlopen("http://127.0.0.1:8188/system_stats", timeout=1.5) as r:
                if r.getcode() == 200:
                    ready = True
                    break
        except (_ue.URLError, OSError, TimeoutError):
            _t.sleep(0.8)
    if ready:
        print("  ComfyUI:    ready ✓", flush=True)
    else:
        print("  ComfyUI:    not responding yet — continuing anyway", flush=True)
    return proc


def open_web_app(url: str, host: str, port: int) -> None:
    """Open the ComfyDock web app in the default browser once uvicorn is up."""
    import time as _t
    import urllib.request as _u
    import urllib.error as _ue

    def _wait_and_open():
        deadline = _t.time() + 15.0
        while _t.time() < deadline:
            try:
                with _u.urlopen(f"http://{host}:{port}/", timeout=1.0) as r:
                    if r.getcode() < 500:
                        break
            except (_ue.URLError, OSError, TimeoutError):
                _t.sleep(0.4)
        try:
            webbrowser.open(url)
            print(f"  Browser:    opened {url}", flush=True)
        except Exception as e:
            print(f"  Browser:    could not open ({e})", flush=True)

    threading.Thread(target=_wait_and_open, daemon=True).start()


def main():
    if args.print_pair_code:
        cfg = _read_config()
        code = cfg.get("cloud_pairing_code")
        if cfg.get("cloud_paired"):
            print("This device is already paired. No pair code needed.")
            sys.exit(0)
        if code:
            print(code)
            sys.exit(0)
        print("No pair code yet — run the agent normally once to register.", file=sys.stderr)
        sys.exit(1)

    print("=" * 60)
    print(f"  ComfyDock Agent v{VERSION}")
    print("=" * 60)
    print(f"  Listening:  http://{args.host}:{args.port}")
    print(f"  Local access token (NOT the pair code): {TOKEN}")
    print(f"              (paste into ComfyDock → Settings → Access token)")
    print(f"  ComfyUI:    {COMFY_URL}")
    print(f"  Comfy path: {COMFY_PATH}")
    if not args.no_cloud:
        start_cloud_relay()
    else:
        print("  Cloud relay: disabled (--no-cloud)")
    if not args.no_comfy:
        proc = launch_comfyui()
        if proc is not None:
            globals()["COMFY_PROC"] = proc
            globals()["_relaunch_comfy"] = launch_comfyui
    else:
        print("  ComfyUI:    not launched (--no-comfy)", flush=True)
    if not args.no_open:
        open_web_app(args.web_url, args.host, args.port)
    else:
        print("  Browser:    auto-open disabled (--no-open)", flush=True)
    print("=" * 60)
    # Re-print pair box at the very end so ComfyUI startup logs don't bury it.
    if _LAST_PAIR_CODE:
        _print_pair_banner(_LAST_PAIR_CODE, _LAST_PAIR_LABEL)
        print(f"  (pair code also saved to: {PAIR_CODE_FILE})", flush=True)
    uvicorn.run(app, host=args.host, port=args.port, log_level="warning")


if __name__ == "__main__":
    try:
        main()
    except SystemExit:
        raise
    except BaseException as _exc:
        import traceback as _tb
        print("", flush=True)
        print("=" * 60, flush=True)
        print("  AGENT CRASHED — full traceback below", flush=True)
        print("=" * 60, flush=True)
        _tb.print_exc()
        print("=" * 60, flush=True)
    finally:
        # Keep the console window open on Windows so the user can read errors
        # / pair code instead of the window vanishing.
        try:
            if sys.stdin and sys.stdin.isatty():
                input("\nPress Enter to close this window…")
        except Exception:
            pass