"""Admin UI — Local FastAPI server to manage ComfyUI Serverless deployment."""

from __future__ import annotations

import ast
import asyncio
import json
import os
import re
import shlex
import signal
import subprocess
import sys
from pathlib import Path
from typing import Any, Optional

shlex_quote = shlex.quote

import httpx
from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

ROOT_DIR = Path(__file__).resolve().parent.parent
# .env lives one level up from ROOT_DIR (at the unzip root), alongside the
# .command/.bat launchers — so users can find and edit it easily.
# We also support a legacy .env inside ROOT_DIR (old layout) as a fallback.
_UNZIP_ROOT = ROOT_DIR.parent
ENV_FILE_ROOT = _UNZIP_ROOT / ".env"
ENV_FILE_LEGACY = ROOT_DIR / ".env"
ENV_FILE = ENV_FILE_ROOT if not ENV_FILE_LEGACY.exists() else ENV_FILE_LEGACY

COMFY_PY = ROOT_DIR / "comfy_server" / "comfy.py"
REQ_TXT = ROOT_DIR / "comfy_server" / "requirements.txt"
DEPLOY_SH = ROOT_DIR / "deploy.sh"
PROXY_PY = ROOT_DIR / "proxy" / "fal_comfy_server.py"
VENV_PY = ROOT_DIR / ".venv" / "bin" / "python"
if not VENV_PY.exists():
    VENV_PY = ROOT_DIR / ".venv" / "Scripts" / "python.exe"

app = FastAPI()

# Proxy subprocess state (single instance — lock prevents race on double-click)
_proxy_proc: Optional[asyncio.subprocess.Process] = None
_proxy_log_buf: list[str] = []
_PROXY_LOG_MAX = 500
_proxy_lock: Optional[asyncio.Lock] = None


def _get_proxy_lock() -> asyncio.Lock:
    global _proxy_lock
    if _proxy_lock is None:
        _proxy_lock = asyncio.Lock()
    return _proxy_lock


# ═══════════════════════════════════════════════════════════════════════════
# comfy.py parser — edit MODELS and CUSTOM_NODES safely via AST
# ═══════════════════════════════════════════════════════════════════════════


def read_comfy_config() -> dict:
    """Return MODELS, CUSTOM_NODES, CUSTOM_NODE_MODELS, COMFY_COMMIT_HASH from comfy.py.

    If comfy.py has a syntax error (e.g. corrupted by a previous bad write),
    return an empty config with an `error` flag so the UI can surface it
    instead of crashing.
    """
    src = COMFY_PY.read_text()
    result = {
        "models": {},
        "custom_nodes": [],
        "custom_node_models": {},
        "comfy_commit_hash": "",
    }
    try:
        tree = ast.parse(src)
    except SyntaxError as e:
        result["error"] = f"comfy.py has a syntax error at line {e.lineno}: {e.msg}"
        return result
    for node in ast.walk(tree):
        # Handle both `X = ...` (Assign) and `X: type = ...` (AnnAssign)
        if isinstance(node, ast.AnnAssign):
            if not isinstance(node.target, ast.Name) or node.value is None:
                continue
            targets = [(node.target.id, node.value)]
        elif isinstance(node, ast.Assign):
            targets = [
                (t.id, node.value) for t in node.targets if isinstance(t, ast.Name)
            ]
        else:
            continue

        for name, value in targets:
            # Only accept literal Dict/List values — skip fallbacks like `MODELS = config['models']`
            if name == "MODELS" and isinstance(value, ast.Dict):
                result["models"] = _eval_dict_of_lists(value)
            elif name == "CUSTOM_NODES" and isinstance(value, ast.List):
                result["custom_nodes"] = _eval_list_of_lists(value)
            elif name == "CUSTOM_NODE_MODELS" and isinstance(value, ast.Dict):
                result["custom_node_models"] = _eval_dict_of_lists(value)
            elif name == "COMFY_COMMIT_HASH" and isinstance(value, ast.Constant):
                result["comfy_commit_hash"] = value.value
    return result


def _eval_dict_of_lists(node: ast.expr) -> dict:
    """Turn a dict literal of {str: [[str, str], ...]} into a Python dict."""
    if not isinstance(node, ast.Dict):
        return {}
    out: dict[str, list[list[str]]] = {}
    for k, v in zip(node.keys, node.values):
        if not isinstance(k, ast.Constant) or not isinstance(k.value, str):
            continue
        out[k.value] = _eval_list_of_lists(v)
    return out


def _eval_list_of_lists(node: ast.expr) -> list[list[str]]:
    if not isinstance(node, ast.List):
        return []
    out: list[list[str]] = []
    for item in node.elts:
        if not isinstance(item, ast.List):
            continue
        pair: list[str] = []
        for sub in item.elts:
            if isinstance(sub, ast.Constant) and isinstance(sub.value, str):
                pair.append(sub.value)
        if len(pair) >= 2:
            out.append(pair[:2])
    return out


def _render_list_of_lists(pairs: list[list[str]], indent: int) -> str:
    pad = " " * indent
    return "\n".join(
        f'{pad}[{json.dumps(p[0])}, {json.dumps(p[1])}],'
        for p in pairs
    )


def _render_dict_of_lists(d: dict[str, list[list[str]]], indent: int) -> str:
    pad = " " * indent
    lines = []
    for folder, pairs in d.items():
        lines.append(f'{pad}{json.dumps(folder)}: [')
        lines.append(_render_list_of_lists(pairs, indent=indent + 4))
        lines.append(f'{pad}],')
    return "\n".join(lines)


def _render_list_literal(items: list, indent: int = 0) -> str:
    """Render a Python list literal. Used for MODELS dict, CUSTOM_NODES list, etc."""
    pad = " " * indent
    if isinstance(items, dict):
        inner = _render_dict_of_lists(items, indent + 4)
        return f"{pad}{{\n{inner}\n{pad}}}"
    else:
        inner = _render_list_of_lists(items, indent + 4)
        return f"{pad}[\n{inner}\n{pad}]"


def _locate_module_assignment(src: str, var_name: str) -> tuple[int, int, int, int] | None:
    """Find the *literal* module-level assignment for `var_name` (AnnAssign or Assign
    whose RHS is a Dict/List literal). Return (start_line, start_col, end_line, end_col)
    of the RHS value — all 1-based line numbers, 0-based columns, as returned by ast.

    Skips reassignments like `MODELS = config["models"]` (value is a Subscript, not a literal).
    """
    tree = ast.parse(src)
    for node in tree.body:  # module-level only, no walk into functions
        target_name = None
        value = None
        if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.value is not None:
            target_name = node.target.id
            value = node.value
        elif isinstance(node, ast.Assign):
            if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
                target_name = node.targets[0].id
                value = node.value
        if target_name != var_name or value is None:
            continue
        if not isinstance(value, (ast.Dict, ast.List)):
            continue  # skip non-literal assignments like `MODELS = config["models"]`
        return (value.lineno, value.col_offset, value.end_lineno, value.end_col_offset)
    return None


def _splice_value(src: str, location: tuple[int, int, int, int], new_literal: str) -> str:
    """Replace the byte range described by `location` with `new_literal`, preserving
    the rest of the file verbatim."""
    start_line, start_col, end_line, end_col = location
    lines = src.splitlines(keepends=True)
    # Build prefix (everything before the value starts)
    prefix = "".join(lines[: start_line - 1]) + lines[start_line - 1][:start_col]
    # Build suffix (everything after the value ends)
    suffix = lines[end_line - 1][end_col:] + "".join(lines[end_line:])
    return prefix + new_literal + suffix


def write_comfy_config(models: dict, custom_nodes: list, custom_node_models: dict):
    """Rewrite comfy.py with new MODELS / CUSTOM_NODES / CUSTOM_NODE_MODELS.

    Uses AST to locate the literal assignments (so regex fragility / unbalanced
    braces can never corrupt the file). After writing, re-parses the result
    and rolls back on syntax error.
    """
    original = COMFY_PY.read_text()

    updates = [
        ("MODELS", _render_list_literal(models, indent=0)),
        ("CUSTOM_NODES", _render_list_literal(custom_nodes, indent=0)),
        ("CUSTOM_NODE_MODELS", _render_list_literal(custom_node_models, indent=0)),
    ]

    src = original
    for var_name, new_literal in updates:
        loc = _locate_module_assignment(src, var_name)
        if loc is None:
            # Variable not found or not a literal — skip silently, never corrupt
            continue
        src = _splice_value(src, loc, new_literal)

    # Validate: the new file must still parse. If not, roll back entirely.
    try:
        ast.parse(src)
    except SyntaxError as e:
        raise HTTPException(
            500,
            f"Refused to save — generated comfy.py would have a syntax error "
            f"({e.msg} at line {e.lineno}). File left untouched.",
        )

    COMFY_PY.write_text(src)


# ═══════════════════════════════════════════════════════════════════════════
# requirements.txt management
# ═══════════════════════════════════════════════════════════════════════════


def read_requirements() -> list[str]:
    if not REQ_TXT.exists():
        return []
    return REQ_TXT.read_text().splitlines()


def write_requirements(lines: list[str]):
    REQ_TXT.write_text("\n".join(lines) + "\n")


# ═══════════════════════════════════════════════════════════════════════════
# Custom node deps fetcher
# ═══════════════════════════════════════════════════════════════════════════


async def fetch_custom_node_requirements(git_url: str, commit: str) -> list[str]:
    """Fetch requirements.txt from a custom node's GitHub repo at a given commit."""
    m = re.match(r"https://github\.com/([^/]+)/([^/.]+)(?:\.git)?/?$", git_url.strip())
    if not m:
        return []
    owner, repo = m.group(1), m.group(2)
    url = f"https://raw.githubusercontent.com/{owner}/{repo}/{commit}/requirements.txt"
    async with httpx.AsyncClient(timeout=10) as client:
        try:
            resp = await client.get(url)
            if resp.status_code != 200:
                return []
            return [ln.strip() for ln in resp.text.splitlines() if ln.strip() and not ln.strip().startswith("#")]
        except Exception:
            return []


# ═══════════════════════════════════════════════════════════════════════════
# fal API — runners
# ═══════════════════════════════════════════════════════════════════════════


def read_env() -> dict:
    if not ENV_FILE.exists():
        return {}
    out = {}
    for ln in ENV_FILE.read_text().splitlines():
        if "=" in ln and not ln.startswith("#"):
            k, v = ln.split("=", 1)
            out[k.strip()] = v.strip()
    return out


def write_env(env: dict):
    """Persist key/value pairs to .env, preserving existing keys not in update."""
    current = read_env()
    current.update({k: v for k, v in env.items() if v is not None})
    lines = [f"{k}={v}" for k, v in current.items() if v != ""]
    ENV_FILE.write_text("\n".join(lines) + "\n")


async def list_runners() -> list[dict]:
    """Return runners scoped to APP_NAME; mark the one serving our local proxy."""
    env = read_env()
    key = env.get("FAL_KEY")
    if not key:
        return []
    app_name = env.get("APP_NAME", "comfy")

    # Is the local proxy running? If yes, figure out which runner is "yours"
    proxy_up = False
    my_runner_id: Optional[str] = None
    try:
        async with httpx.AsyncClient(timeout=2) as c:
            r = await c.get("http://localhost:8188/api/my_runner")
            if r.status_code == 200:
                proxy_up = True
                my_runner_id = r.json().get("runner_id")
    except Exception:
        pass

    async with httpx.AsyncClient(timeout=10) as client:
        try:
            resp = await client.get(
                "https://rest.alpha.fal.ai/runners/",
                headers={"Authorization": f"Key {key}"},
            )
            data = resp.json()
            runners = data if isinstance(data, list) else data.get("runners", [])
            out = []
            for r in runners:
                if r.get("application") != app_name:
                    continue
                r["is_mine"] = bool(my_runner_id and r.get("runner_id") == my_runner_id)
                out.append(r)

            # Only mark "YOURS" when the proxy has actually captured the
            # x-fal-runner-id from a real prompt response. No heuristics — we
            # don't want to mislabel another user's runner as yours.

            # Put the user's runner first
            out.sort(key=lambda x: (not x.get("is_mine"), x.get("started_at", "")))
            return out
        except Exception:
            return []


# ═══════════════════════════════════════════════════════════════════════════
# API endpoints
# ═══════════════════════════════════════════════════════════════════════════


@app.get("/api/config")
async def get_config():
    cfg = read_comfy_config()
    cfg["requirements"] = read_requirements()
    env = read_env()
    # Mask the API key before sending to the browser
    safe_env = dict(env)
    if safe_env.get("FAL_KEY"):
        key = safe_env["FAL_KEY"]
        safe_env["FAL_KEY_MASKED"] = key[:6] + "…" + key[-4:] if len(key) > 12 else "●●●●"
        safe_env.pop("FAL_KEY")
    # Mark token presence without leaking the value to the browser
    if safe_env.get("HF_TOKEN"):
        safe_env["HF_TOKEN_SET"] = True
        safe_env.pop("HF_TOKEN")
    if safe_env.get("CIVITAI_TOKEN"):
        safe_env["CIVITAI_TOKEN_SET"] = True
        safe_env.pop("CIVITAI_TOKEN")
    cfg["env"] = safe_env
    return cfg


@app.get("/api/env_status")
async def env_status():
    env = read_env()
    return {
        "has_key": bool(env.get("FAL_KEY")),
        "app_name": env.get("APP_NAME", "comfy"),
    }


@app.get("/api/scaling")
async def get_scaling():
    env = read_env()
    key = env.get("FAL_KEY")
    app_name = env.get("APP_NAME", "comfy")
    if not key:
        return {"ok": False, "error": "no_key"}
    async with httpx.AsyncClient(timeout=8) as client:
        try:
            resp = await client.get(
                "https://rest.alpha.fal.ai/applications/?limit=50",
                headers={"Authorization": f"Key {key}"},
            )
            if resp.status_code != 200:
                return {"ok": False, "error": f"fal_{resp.status_code}"}
            apps = resp.json()
            apps = apps if isinstance(apps, list) else apps.get("applications", [])
            for a in apps:
                if a.get("app_name") == app_name:
                    return {
                        "ok": True,
                        "min_concurrency": a.get("min_concurrency"),
                        "max_concurrency": a.get("max_concurrency"),
                        "max_multiplexing": a.get("max_multiplexing"),
                        "keep_alive": a.get("keep_alive"),
                    }
            return {"ok": False, "error": "app_not_found"}
        except Exception as e:
            return {"ok": False, "error": str(e)}


@app.post("/api/scaling")
async def set_scaling(request: Request):
    """Run `fal apps scale` with provided fields."""
    body = await request.json()
    env = read_env()
    if not env.get("FAL_KEY"):
        raise HTTPException(400, "FAL_KEY missing")
    app_name = env.get("APP_NAME", "comfy")

    flags: list[str] = []
    for key, flag in [
        ("min_concurrency", "--min-concurrency"),
        ("max_concurrency", "--max-concurrency"),
        ("max_multiplexing", "--max-multiplexing"),
        ("keep_alive", "--keep-alive"),
    ]:
        v = body.get(key)
        if v is not None and str(v).strip() != "":
            flags += [flag, str(int(v))]
    if not flags:
        raise HTTPException(400, "Nothing to update")

    if IS_WINDOWS:
        # Windows has no bash; run fal CLI directly via VENV_PY with env loaded.
        proc = await asyncio.create_subprocess_exec(
            str(VENV_PY), "-m", "fal", "apps", "scale", app_name, *flags,
            cwd=str(ROOT_DIR),
            env={**os.environ, **read_env()},
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
    else:
        proc = await asyncio.create_subprocess_exec(
            "bash", "-c",
            f"cd {shlex_quote(str(ROOT_DIR))} && "
            ".venv/bin/fal auth logout >/dev/null 2>&1 || true; "
            f"set -a; source {shlex_quote(str(ENV_FILE))}; set +a; "
            f".venv/bin/fal apps scale {app_name} " + " ".join(flags) + " 2>&1",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
    out, _ = await proc.communicate()
    text = out.decode(errors="replace")
    return {"ok": proc.returncode == 0, "output": text}


@app.get("/api/whoami")
async def whoami():
    """Return the fal account/team the current FAL_KEY belongs to."""
    env = read_env()
    key = env.get("FAL_KEY")
    if not key:
        return {"ok": False, "error": "no_key"}
    async with httpx.AsyncClient(timeout=8) as client:
        try:
            resp = await client.get(
                "https://rest.alpha.fal.ai/applications/?limit=1",
                headers={"Authorization": f"Key {key}"},
            )
            if resp.status_code == 401:
                return {"ok": False, "error": "invalid_key"}
            if resp.status_code != 200:
                return {"ok": False, "error": f"fal_api_{resp.status_code}"}
            data = resp.json()
            apps = data if isinstance(data, list) else data.get("applications", [])
            if not apps:
                # No apps yet — key works but no apps deployed
                return {"ok": True, "slug": "(no apps yet)", "has_apps": False}
            return {
                "ok": True,
                "slug": apps[0].get("app_user_nickname") or "unknown",
                "has_apps": True,
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}


@app.get("/api/playground_url")
async def playground_url():
    """Resolve the fal.ai playground URL for the currently configured app."""
    env = read_env()
    key = env.get("FAL_KEY")
    app_name = env.get("APP_NAME", "comfy")
    if not key:
        return {"ok": False, "error": "no_key"}
    async with httpx.AsyncClient(timeout=10) as client:
        try:
            resp = await client.get(
                "https://rest.alpha.fal.ai/applications/?limit=100",
                headers={"Authorization": f"Key {key}"},
            )
            if resp.status_code != 200:
                return {"ok": False, "error": f"fal api {resp.status_code}"}
            data = resp.json()
            apps = data if isinstance(data, list) else data.get("applications", [])
            slug = None
            for a in apps:
                if a.get("app_name") == app_name:
                    slug = a.get("app_user_nickname")
                    break
            if not slug and apps:
                slug = apps[0].get("app_user_nickname")
            if not slug:
                return {"ok": False, "error": "no app found on this account"}
            return {
                "ok": True,
                "url": f"https://fal.ai/models/{slug}/{app_name}",
                "slug": slug,
                "app": app_name,
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}


@app.post("/api/env")
async def save_env(request: Request):
    body = await request.json()
    updates = {}
    for k in ("FAL_KEY", "APP_NAME", "FAL_COMFY_ENDPOINT", "HF_TOKEN", "CIVITAI_TOKEN"):
        if k in body and body[k] is not None:
            updates[k] = str(body[k]).strip()
    if not updates:
        raise HTTPException(400, "No known keys provided")
    write_env(updates)
    return {"ok": True}


@app.post("/api/models")
async def save_models(request: Request):
    body = await request.json()
    cfg = read_comfy_config()
    write_comfy_config(body["models"], cfg["custom_nodes"], cfg["custom_node_models"])
    return {"ok": True}


@app.post("/api/custom_nodes")
async def save_custom_nodes(request: Request):
    body = await request.json()
    cfg = read_comfy_config()
    write_comfy_config(cfg["models"], body["custom_nodes"], cfg["custom_node_models"])
    return {"ok": True}


def write_comfy_commit(commit: str):
    """Replace the COMFY_COMMIT_HASH assignment in comfy.py."""
    src = COMFY_PY.read_text()
    # Match `COMFY_COMMIT_HASH = "..."` or with annotation
    pattern = re.compile(
        r'^(COMFY_COMMIT_HASH\s*(?::\s*[^=]+)?\s*=\s*)(["\'])[^"\']*(["\'])',
        re.MULTILINE,
    )
    new_src, n = pattern.subn(rf'\g<1>"{commit}"', src, count=1)
    if n == 0:
        raise HTTPException(400, "COMFY_COMMIT_HASH not found in comfy.py")
    COMFY_PY.write_text(new_src)


@app.post("/api/commit_hash")
async def save_commit_hash(request: Request):
    body = await request.json()
    commit = (body.get("commit") or "").strip()
    if not re.fullmatch(r"[0-9a-f]{7,40}", commit):
        raise HTTPException(400, "Not a valid git SHA (7–40 hex chars)")
    # Validate against upstream ComfyUI
    async with httpx.AsyncClient(timeout=10) as client:
        resp = await client.get(
            f"https://api.github.com/repos/comfyanonymous/ComfyUI/commits/{commit}"
        )
        if resp.status_code != 200:
            raise HTTPException(
                400,
                f"Commit not found on comfyanonymous/ComfyUI (HTTP {resp.status_code})",
            )
        data = resp.json()
        full_sha = data.get("sha", commit)
        message = (data.get("commit") or {}).get("message", "").split("\n")[0]
    write_comfy_commit(full_sha)
    return {"ok": True, "commit": full_sha, "message": message}


@app.post("/api/requirements")
async def save_requirements(request: Request):
    body = await request.json()
    write_requirements(body["requirements"])
    return {"ok": True}


@app.post("/api/fetch_node_deps")
async def fetch_node_deps(request: Request):
    body = await request.json()
    deps = await fetch_custom_node_requirements(body["git_url"], body["commit"])
    return {"deps": deps}


@app.get("/api/runners")
async def get_runners():
    return await list_runners()


ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")


@app.get("/api/runner_logs/{runner_id}")
async def runner_logs(runner_id: str):
    """Poll the fal REST logs API and stream new lines back to the browser."""
    env = read_env()
    key = env.get("FAL_KEY")
    if not key:
        return JSONResponse({"error": "No FAL_KEY"}, status_code=400)

    async def stream():
        seen: set[str] = set()
        yield f"[connecting to fal logs for {runner_id[:12]}…]\n".encode()
        headers = {"Authorization": f"Key {key}"}
        url = f"https://rest.alpha.fal.ai/logs/?runner_id={runner_id}&limit=200"
        async with httpx.AsyncClient(timeout=10, headers=headers) as client:
            # Initial backfill + then poll for new entries
            while True:
                try:
                    resp = await client.get(url)
                    if resp.status_code != 200:
                        yield f"[fal api error {resp.status_code}] {resp.text[:200]}\n".encode()
                        await asyncio.sleep(3)
                        continue
                    entries = resp.json()
                    # API returns newest-first; we want oldest-first on screen
                    entries_sorted = sorted(entries, key=lambda e: e.get("timestamp", ""))
                    new_lines = []
                    for e in entries_sorted:
                        sig = f"{e.get('timestamp','')}|{e.get('message','')[:80]}"
                        if sig in seen:
                            continue
                        seen.add(sig)
                        ts = (e.get("timestamp") or "")[11:19]
                        lvl = (e.get("level") or "INFO").ljust(5)
                        msg = ANSI_RE.sub("", e.get("message") or "")
                        new_lines.append(f"{ts} {lvl} {msg}")
                    if new_lines:
                        yield ("\n".join(new_lines) + "\n").encode()
                except Exception as e:
                    yield f"[stream error: {e}]\n".encode()
                await asyncio.sleep(2)

    return StreamingResponse(stream(), media_type="text/plain")


@app.post("/api/delete_app")
async def delete_app(request: Request):
    """Undeploy the current app via `fal apps delete <app>`. Requires explicit
    confirmation in the body ({"confirm": "<app_name>"}) so it can't be
    triggered by accident."""
    body = await request.json()
    env = read_env()
    if not env.get("FAL_KEY"):
        raise HTTPException(400, "FAL_KEY missing")
    app_name = env.get("APP_NAME", "comfy")
    if body.get("confirm") != app_name:
        raise HTTPException(400, f"Type the app name ({app_name}) to confirm deletion")

    if IS_WINDOWS:
        # Windows has no bash and no `yes`; pipe a stream of "y\n" via stdin instead.
        proc = await asyncio.create_subprocess_exec(
            str(VENV_PY), "-m", "fal", "apps", "delete", app_name,
            cwd=str(ROOT_DIR),
            env={**os.environ, **read_env()},
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
        out, _ = await proc.communicate(input=b"y\n" * 5)
    else:
        proc = await asyncio.create_subprocess_exec(
            "bash", "-c",
            f"cd {shlex_quote(str(ROOT_DIR))} && "
            ".venv/bin/fal auth logout >/dev/null 2>&1 || true; "
            f"set -a; source {shlex_quote(str(ENV_FILE))}; set +a; "
            f"yes | .venv/bin/fal apps delete {app_name} 2>&1",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
        out, _ = await proc.communicate()
    return {
        "ok": proc.returncode == 0,
        "app": app_name,
        "output": out.decode(errors="replace"),
    }


@app.post("/api/kill_runner/{runner_id}")
async def kill_runner(runner_id: str):
    env = read_env()
    key = env.get("FAL_KEY")
    if not key:
        return {"ok": False, "error": "No FAL_KEY"}
    # Use the fal CLI
    try:
        result = subprocess.run(
            ["fal", "runners", "kill", runner_id],
            env={**os.environ, "FAL_KEY": key},
            capture_output=True,
            text=True,
            timeout=10,
        )
        return {"ok": result.returncode == 0, "output": result.stdout + result.stderr}
    except Exception as e:
        return {"ok": False, "error": str(e)}


@app.get("/api/validate_url")
async def validate_url(url: str):
    """Check if a URL is reachable. Uses HF_TOKEN/CIVITAI_TOKEN from .env when
    the URL points at a gated host, so the validator doesn't false-positive
    on a 401 that the runner would actually pass through."""
    env = read_env()
    headers = {}
    if "huggingface.co" in url and env.get("HF_TOKEN"):
        headers["Authorization"] = f"Bearer {env['HF_TOKEN']}"
    elif "civitai.com" in url and env.get("CIVITAI_TOKEN"):
        headers["Authorization"] = f"Bearer {env['CIVITAI_TOKEN']}"

    async with httpx.AsyncClient(timeout=10, follow_redirects=True, headers=headers) as client:
        try:
            # HEAD is sometimes blocked → fall back to a tiny GET range
            resp = await client.head(url)
            if resp.status_code in (405, 403):
                resp = await client.get(url, headers={**headers, "Range": "bytes=0-0"})
            return {"ok": resp.status_code in (200, 206), "status": resp.status_code}
        except Exception as e:
            return {"ok": False, "error": str(e)}


@app.get("/api/validate_commit")
async def validate_commit(git_url: str, commit: str):
    """Check if a commit exists in a GitHub repo."""
    m = re.match(r"https://github\.com/([^/]+)/([^/.]+)(?:\.git)?/?$", git_url.strip())
    if not m:
        return {"ok": False, "error": "Invalid GitHub URL"}
    owner, repo = m.group(1), m.group(2)
    url = f"https://api.github.com/repos/{owner}/{repo}/commits/{commit}"
    async with httpx.AsyncClient(timeout=10) as client:
        try:
            resp = await client.get(url)
            return {"ok": resp.status_code == 200}
        except Exception as e:
            return {"ok": False, "error": str(e)}


@app.post("/api/deploy")
async def deploy():
    """Stream deploy.sh output."""
    async def stream():
        # Non-interactive: expects .env to be set already
        env = {**os.environ, **read_env()}
        env_vars = read_env()
        app_name = env_vars.get("APP_NAME", "comfy")

        if IS_WINDOWS:
            # Windows has no bash. Run fal CLI directly through VENV_PY.
            # Push tokens as separate subprocess calls before the deploy.
            for token_name in ("HF_TOKEN", "CIVITAI_TOKEN"):
                token_val = env_vars.get(token_name)
                if token_val:
                    yield f"Pushing {token_name} to fal secrets...\n".encode()
                    secret_proc = await asyncio.create_subprocess_exec(
                        str(VENV_PY), "-m", "fal", "secrets", "set",
                        f"{token_name}={token_val}",
                        cwd=str(ROOT_DIR),
                        env=env,
                        stdout=asyncio.subprocess.PIPE,
                        stderr=asyncio.subprocess.STDOUT,
                    )
                    s_out, _ = await secret_proc.communicate()
                    if s_out:
                        yield s_out

            proc = await asyncio.create_subprocess_exec(
                str(VENV_PY), "-m", "fal", "deploy",
                "comfy_server/server.py::Comfy",
                "--app-name", app_name,
                "--auth", "private",
                cwd=str(ROOT_DIR),
                env=env,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
            )
        else:
            # Push tokens to fal account-level secrets BEFORE deploy so the
            # runner can read them via os.getenv at runtime. The runner's
            # download_utils reads HF_TOKEN / CIVITAI_TOKEN to authenticate
            # gated downloads (e.g. Black Forest Labs Flux 2 Klein).
            token_lines = []
            for token_name in ("HF_TOKEN", "CIVITAI_TOKEN"):
                token_val = env_vars.get(token_name)
                if token_val:
                    # fal secrets set syntax: NAME=VALUE (no space)
                    token_lines.append(
                        f"echo 'Pushing {token_name} to fal secrets...' && "
                        f".venv/bin/python -m fal secrets set \"{token_name}=${token_name}\" 2>&1 || true; "
                    )
            push_secrets = "".join(token_lines)

            proc = await asyncio.create_subprocess_exec(
                "bash", "-c",
                f"cd {shlex_quote(str(ROOT_DIR))} && "
                ".venv/bin/fal auth logout >/dev/null 2>&1 || true; "
                "set -a; source " + shlex_quote(str(ENV_FILE)) + "; set +a; "
                + push_secrets +
                ".venv/bin/python -m fal deploy comfy_server/server.py::Comfy --app-name \"${APP_NAME:-comfy}\" --auth private 2>&1",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
                env=env,
            )
        assert proc.stdout
        async for line in proc.stdout:
            yield line
        await proc.wait()
        if proc.returncode == 0:
            yield b"\n[DEPLOY SUCCESSFUL] Your app is live on fal.ai.\n"
        else:
            yield f"\n[DEPLOY FAILED — exit code {proc.returncode}] Check the log above.\n".encode()
        # Keep a parseable sentinel at the very end so the frontend can detect status
        yield f"[exit code: {proc.returncode}]\n".encode()

    return StreamingResponse(stream(), media_type="text/plain")


# ═══════════════════════════════════════════════════════════════════════════
# Local ComfyUI proxy — start/stop/logs
# ═══════════════════════════════════════════════════════════════════════════


async def _drain_proxy_logs():
    global _proxy_proc
    if not _proxy_proc or not _proxy_proc.stdout:
        return
    async for line in _proxy_proc.stdout:
        text = line.decode(errors="replace").rstrip()
        _proxy_log_buf.append(text)
        if len(_proxy_log_buf) > _PROXY_LOG_MAX:
            del _proxy_log_buf[: len(_proxy_log_buf) - _PROXY_LOG_MAX]


@app.get("/api/proxy/status")
async def proxy_status():
    running = _proxy_proc is not None and _proxy_proc.returncode is None
    # Opportunistic peek at runner state (for UI hint like "runner IDLE / STARTING")
    runner_state = None
    if running:
        env = read_env()
        key = env.get("FAL_KEY")
        if key:
            try:
                async with httpx.AsyncClient(timeout=4) as client:
                    resp = await client.get(
                        "https://rest.alpha.fal.ai/runners/",
                        headers={"Authorization": f"Key {key}"},
                    )
                    if resp.status_code == 200:
                        data = resp.json()
                        runners = data if isinstance(data, list) else data.get("runners", [])
                        # Pick the most recent runner for our app
                        app_name = env.get("APP_NAME", "comfy")
                        ours = [r for r in runners if r.get("application") == app_name]
                        if ours:
                            runner_state = ours[0].get("state")
            except Exception:
                pass
    return {
        "running": running,
        "url": "http://localhost:8188" if running else "",
        "runner_state": runner_state,
    }


@app.post("/api/proxy/start")
async def proxy_start():
    global _proxy_proc
    async with _get_proxy_lock():
        return await _proxy_start_locked()


IS_WINDOWS = sys.platform.startswith("win") or os.name == "nt"


def _find_pids_on_port(port: int) -> list[int]:
    """Return PIDs of processes listening on `port`, cross-platform."""
    pids: set[int] = set()
    try:
        if IS_WINDOWS:
            # netstat -ano | findstr ":PORT" — parse PID from last column
            result = subprocess.run(
                ["netstat", "-ano"], capture_output=True, text=True, timeout=5,
            )
            for line in result.stdout.splitlines():
                if f":{port}" not in line or "LISTENING" not in line:
                    continue
                parts = line.split()
                if parts and parts[-1].isdigit():
                    pids.add(int(parts[-1]))
        else:
            result = subprocess.run(
                ["lsof", "-ti", f":{port}"], capture_output=True, text=True, timeout=3,
            )
            for line in result.stdout.strip().splitlines():
                line = line.strip()
                if line.isdigit():
                    pids.add(int(line))
    except Exception:
        pass
    return list(pids)


def _kill_port(port: int):
    """Kill any process bound to `port` (orphans from crashed sessions)."""
    my_pid = os.getpid()
    for pid in _find_pids_on_port(port):
        if pid == my_pid:
            continue
        try:
            if IS_WINDOWS:
                subprocess.run(
                    ["taskkill", "/F", "/PID", str(pid)],
                    capture_output=True, timeout=3,
                )
            else:
                os.kill(pid, signal.SIGTERM)
        except (ProcessLookupError, ValueError, subprocess.SubprocessError):
            pass

    # Sweep stale proxy processes by name (Unix only — Windows relies on port kill above)
    if not IS_WINDOWS:
        try:
            subprocess.run(
                ["pkill", "-f", "proxy/fal_comfy_server.py"],
                capture_output=True, timeout=3,
            )
        except Exception:
            pass


async def _proxy_start_locked():
    global _proxy_proc
    if _proxy_proc and _proxy_proc.returncode is None:
        return {"ok": True, "already_running": True}

    # Sweep orphan proxies from crashed sessions before spawning a fresh one
    _kill_port(8188)
    await asyncio.sleep(0.5)
    if not PROXY_PY.exists():
        return JSONResponse(
            {"error": f"Proxy not bundled (missing {PROXY_PY})"}, status_code=500
        )
    env = read_env()
    if not env.get("FAL_KEY"):
        return JSONResponse({"error": "FAL_KEY missing — set it in Settings"}, status_code=400)
    # Use the deployed app's endpoint if not explicitly set
    if not env.get("FAL_COMFY_ENDPOINT"):
        # Resolve from whoami + APP_NAME
        app_name = env.get("APP_NAME", "comfy")
        try:
            async with httpx.AsyncClient(timeout=8) as client:
                resp = await client.get(
                    "https://rest.alpha.fal.ai/applications/?limit=1",
                    headers={"Authorization": f"Key {env['FAL_KEY']}"},
                )
                apps = resp.json() if resp.status_code == 200 else []
                if isinstance(apps, list) and apps:
                    slug = apps[0].get("app_user_nickname")
                    if slug:
                        env["FAL_COMFY_ENDPOINT"] = f"{slug}/{app_name}"
        except Exception:
            pass
    if not env.get("FAL_COMFY_ENDPOINT"):
        return JSONResponse(
            {"error": "Could not resolve FAL_COMFY_ENDPOINT — deploy first or set it manually"},
            status_code=400,
        )

    _proxy_log_buf.clear()
    _proxy_log_buf.append(f"> Starting ComfyUI proxy")
    _proxy_log_buf.append(f"  endpoint : {env['FAL_COMFY_ENDPOINT']}")
    _proxy_log_buf.append(f"  working  : {ROOT_DIR}")
    _proxy_log_buf.append(f"  python   : {VENV_PY}")
    _proxy_log_buf.append("")
    _proxy_proc = await asyncio.create_subprocess_exec(
        str(VENV_PY),
        "-u",  # unbuffered stdout/stderr so logs appear live
        str(PROXY_PY),
        cwd=str(ROOT_DIR),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
        env={**os.environ, **env, "PYTHONUNBUFFERED": "1"},
    )
    asyncio.create_task(_drain_proxy_logs())
    return {"ok": True, "endpoint": env["FAL_COMFY_ENDPOINT"]}


@app.post("/api/proxy/stop")
async def proxy_stop():
    global _proxy_proc
    if not _proxy_proc or _proxy_proc.returncode is not None:
        return {"ok": True, "already_stopped": True}
    try:
        # Windows child processes don't respond to SIGINT the same way as Unix.
        # `terminate()` maps to TerminateProcess on Windows (hard but reliable)
        # and SIGTERM on Unix (graceful).
        if IS_WINDOWS:
            _proxy_proc.terminate()
        else:
            _proxy_proc.send_signal(signal.SIGINT)
        try:
            await asyncio.wait_for(_proxy_proc.wait(), timeout=5)
        except asyncio.TimeoutError:
            _proxy_proc.kill()
            await _proxy_proc.wait()
    finally:
        _proxy_log_buf.append("Proxy stopped.")
    return {"ok": True}


@app.get("/api/proxy/logs")
async def proxy_logs():
    return {"lines": _proxy_log_buf[-200:]}


@app.get("/api/download")
async def download_proxy(url: str, filename: str = "output"):
    """Forward downloads through the local proxy (keeps the force-download behavior)."""
    try:
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.get(
                "http://localhost:8188/api/download",
                params={"url": url, "filename": filename},
            )
            return Response(
                content=resp.content,
                media_type=resp.headers.get("content-type", "application/octet-stream"),
                headers={"Content-Disposition": f'attachment; filename="{filename}"'},
            )
    except Exception as e:
        raise HTTPException(502, f"Proxy unreachable: {e}")


@app.get("/api/download_all")
async def download_all():
    """Stream a ZIP of every generated output from the running proxy."""
    import io
    import zipfile
    from datetime import datetime

    # Fetch outputs list from proxy
    try:
        async with httpx.AsyncClient(timeout=5) as client:
            r = await client.get("http://localhost:8188/api/outputs")
            if r.status_code != 200:
                raise HTTPException(502, f"proxy returned {r.status_code}")
            outputs = r.json()
    except httpx.ConnectError:
        raise HTTPException(503, "ComfyUI proxy not running — start it first")

    # Collect (filename, url) pairs with dedup
    urls: list[tuple[str, str]] = []
    seen_names: set[str] = set()
    for entry in outputs:
        for url in entry.get("urls", []):
            base = url.rsplit("/", 1)[-1] if "/" in url else "output"
            name = base
            i = 2
            while name in seen_names:
                stem, _, ext = base.rpartition(".")
                name = f"{stem}_{i}.{ext}" if ext else f"{base}_{i}"
                i += 1
            seen_names.add(name)
            urls.append((name, url))

    if not urls:
        raise HTTPException(404, "No outputs to download")

    async def stream_zip():
        # Stream the zip by building it in-memory and yielding chunks.
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED, compresslevel=1) as zf:
            async with httpx.AsyncClient(timeout=180) as client:
                for name, url in urls:
                    try:
                        resp = await client.get(url)
                        if resp.status_code == 200:
                            zf.writestr(name, resp.content)
                    except Exception:
                        continue
        buf.seek(0)
        while True:
            chunk = buf.read(64 * 1024)
            if not chunk:
                break
            yield chunk

    ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    return StreamingResponse(
        stream_zip(),
        media_type="application/zip",
        headers={"Content-Disposition": f'attachment; filename="comfyui_outputs_{ts}.zip"'},
    )


@app.get("/api/outputs")
async def outputs_proxy():
    """Forward /api/outputs from the local ComfyUI proxy and deduplicate filenames
    so two outputs with the same basename (e.g. VideoCombine -> `MyVideo_00001.mp4`)
    don't clobber each other on download."""
    try:
        async with httpx.AsyncClient(timeout=5) as client:
            resp = await client.get("http://localhost:8188/api/outputs")
            if resp.status_code != 200:
                return {"ok": False, "error": f"proxy_{resp.status_code}", "outputs": []}
            raw = resp.json()
    except httpx.ConnectError:
        return {"ok": False, "error": "proxy_not_running", "outputs": []}
    except Exception as e:
        return {"ok": False, "error": str(e), "outputs": []}

    # Oldest first -> newest, so older files keep their original name;
    # collisions on later files get _2, _3, … suffix.
    outputs_chrono = sorted(raw, key=lambda e: e.get("timestamp", ""))
    seen: dict[str, int] = {}
    # Walk chrono, annotate each URL with a unique `filename`
    for entry in outputs_chrono:
        entry_files = []
        for url in entry.get("urls", []):
            base = url.rsplit("/", 1)[-1] if "/" in url else "output"
            if base not in seen:
                seen[base] = 1
                name = base
            else:
                seen[base] += 1
                stem, _, ext = base.rpartition(".")
                if stem and ext:
                    name = f"{stem}_{seen[base]}.{ext}"
                else:
                    name = f"{base}_{seen[base]}"
                # guard against the suffixed name already existing
                while name in seen:
                    seen[base] += 1
                    name = f"{stem}_{seen[base]}.{ext}" if stem and ext else f"{base}_{seen[base]}"
                seen[name] = 1
            entry_files.append({"url": url, "filename": name})
        entry["files"] = entry_files
    return {"ok": True, "outputs": outputs_chrono}


# ═══════════════════════════════════════════════════════════════════════════
# Static UI
# ═══════════════════════════════════════════════════════════════════════════


UI_DIR = Path(__file__).parent


@app.get("/")
async def index():
    return FileResponse(UI_DIR / "index.html")


@app.get("/app.js")
async def app_js():
    return FileResponse(UI_DIR / "app.js")


@app.get("/style.css")
async def style_css():
    return FileResponse(UI_DIR / "style.css")


if __name__ == "__main__":
    import atexit
    import webbrowser
    import time
    import threading

    # Write PID file so the launcher can stop only THIS session (not another user's)
    pid_file = ROOT_DIR / ".admin-ui.pid"
    pid_file.write_text(str(os.getpid()))
    atexit.register(lambda: pid_file.unlink(missing_ok=True))

    # Sweep orphan proxies from previous sessions
    _kill_port(8188)

    # Open browser after 1.5s so server has time to start
    def open_browser():
        time.sleep(1.5)
        webbrowser.open("http://localhost:9999")

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

    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=9999, log_level="warning")
