# /// script
# requires-python = ">=3.11"
# dependencies = [
# "fastapi",
# "uvicorn",
# "aiofiles",
# "Pillow",
# "fal-client",
# "python-multipart",
# "sse_starlette",
# "httpx",
# "boto3",
# ]
# ///

import asyncio
import ipaddress
import json
import logging
import os
import re
import sys
import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime
from io import BytesIO
from typing import Any, Dict, List, Optional, Union
from pathlib import Path
from urllib.parse import quote as _urlquote, urlparse

import aiofiles
import fal_client
from fal_client.client import RUN_URL_FORMAT, aconnect_sse
from fastapi import (
    Depends,
    FastAPI,
    File,
    Form,
    Header,
    HTTPException,
    Query,
    Request,
    Response,
    UploadFile,
    WebSocket,
    WebSocketDisconnect,
)
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from PIL import Image
from pydantic import BaseModel, Field
from sse_starlette import EventSourceResponse, ServerSentEvent

FAL_KEY = os.environ.get("FAL_KEY")
if not FAL_KEY:
    print("FAL_KEY environment variable is required")
    sys.exit(1)

ENDPOINT = os.environ.get("FAL_COMFY_ENDPOINT", "comfyui/server")

# --- S3 Configuration (optional) ---
S3_BUCKET = os.environ.get("S3_BUCKET")
S3_REGION = os.environ.get("S3_REGION", "us-east-1")
S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL")  # For Wasabi/MinIO/etc.
S3_PREFIX = os.environ.get("S3_PREFIX", "comfyui")  # Key prefix in the bucket

s3_client = None
if S3_BUCKET:
    import boto3
    s3_kwargs = {"region_name": S3_REGION}
    if S3_ENDPOINT_URL:
        s3_kwargs["endpoint_url"] = S3_ENDPOINT_URL
    s3_client = boto3.client("s3", **s3_kwargs)
    print(f"S3 enabled: bucket={S3_BUCKET}, prefix={S3_PREFIX}")
else:
    print("S3 disabled (set S3_BUCKET to enable)")


async def s3_upload(data: bytes, filename: str, content_type: str, folder: str = "outputs") -> Optional[str]:
    """Upload bytes to S3 and return the key. Returns None if S3 is not configured."""
    if not s3_client:
        return None
    key = f"{S3_PREFIX}/{folder}/{filename}"
    try:
        s3_client.put_object(
            Bucket=S3_BUCKET,
            Key=key,
            Body=data,
            ContentType=content_type,
        )
        print(f"S3: uploaded {key}")
        return key
    except Exception as e:
        print(f"S3 upload failed for {key}: {e}")
        return None


def s3_get_url(key: str) -> str:
    """Generate a presigned URL for an S3 object (valid 1 hour)."""
    params = {"Bucket": S3_BUCKET, "Key": key}
    return s3_client.generate_presigned_url("get_object", Params=params, ExpiresIn=3600)


def s3_upload_sync(data: bytes, filename: str, content_type: str, folder: str = "outputs") -> Optional[str]:
    """Synchronous S3 upload (for use in thread executor)."""
    if not s3_client:
        return None
    key = f"{S3_PREFIX}/{folder}/{filename}"
    try:
        s3_client.put_object(
            Bucket=S3_BUCKET,
            Key=key,
            Body=data,
            ContentType=content_type,
        )
        print(f"S3: uploaded {key}")
        return key
    except Exception as e:
        print(f"S3 upload failed for {key}: {e}")
        return None

def _node_cache_path(endpoint: str) -> str:
    safe_endpoint = re.sub(r"[^A-Za-z0-9._-]+", "_", endpoint)
    return f"node_class_mappings.{safe_endpoint}.json"


def _load_node_class_mappings(endpoint: str) -> dict[str, Any]:
    cache_path = _node_cache_path(endpoint)
    fallback_paths = [
        cache_path,
        "node_class_mappings.comfyui_server.json",
        "node_class_mappings.json",
    ]

    if os.path.exists(cache_path):
        print(f"Using cached node metadata: {cache_path}")
        with open(cache_path, "r") as f:
            return json.load(f)

    org_app = endpoint.replace("/", "_")
    org = endpoint.split("/")[0] if "/" in endpoint else endpoint
    app_name = endpoint.split("/")[1] if "/" in endpoint else ""
    print("")
    print(f"  Fetching node metadata for endpoint: {endpoint}")
    print(f"  -> This starts a GPU runner on fal.ai (cold start takes 3-5 min)")
    print(f"  -> Watch the runner boot in real time at:")
    print(f"    https://fal.ai/dashboard/apps/{org}/{app_name}/runners")
    print("")
    try:
        data = fal_client.subscribe(
            endpoint,
            path="object_info",
            arguments={},
        )
        with open(cache_path, "w") as f:
            json.dump(data, f)
        return data
    except Exception as exc:
        print(f"Failed to fetch node metadata for {endpoint}: {exc}")
        for fallback_path in fallback_paths[1:]:
            if os.path.exists(fallback_path):
                print(f"Falling back to cached node metadata: {fallback_path}")
                with open(fallback_path, "r") as f:
                    return json.load(f)
        print("No cached node metadata available; starting with empty mappings.")
        return {}


def _get_model_choices(node_class: str, input_name: str) -> list[str]:
    try:
        return NODE_CLASS_MAPPINGS[node_class]["input"]["required"][input_name][0]
    except KeyError:
        return []


NODE_CLASS_MAPPINGS = _load_node_class_mappings(ENDPOINT)

MODELS = {
    "checkpoints": _get_model_choices("CheckpointLoaderSimple", "ckpt_name"),
    "loras": _get_model_choices("LoraLoader", "lora_name"),
    "controlnets": _get_model_choices("ControlNetLoader", "control_net_name"),
    "vaes": _get_model_choices("VAELoader", "vae_name"),
    "vae": _get_model_choices("VAELoader", "vae_name"),
    "diffusion_models": _get_model_choices("UNETLoader", "unet_name"),
    "clip": _get_model_choices("CLIPLoader", "clip_name"),
    "text_encoders": _get_model_choices("CLIPLoader", "clip_name"),
    "upscale_models": _get_model_choices("UpscaleModelLoader", "model_name"),
    "style_models": _get_model_choices("StyleModelLoader", "style_model_name"),
}

for model_folder, model_names in MODELS.items():
    new_model_names = []
    for index, model_name in enumerate(model_names):
        new_model_names.append({"name": model_name, "path_index": index})
    MODELS[model_folder] = new_model_names

app = FastAPI()
app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=5)

# Serve the built frontend from dist/
FRONTEND_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "frontend", "dist")
if os.path.isdir(FRONTEND_DIR):
    print(f"Serving frontend from {FRONTEND_DIR}")
    from fastapi.staticfiles import StaticFiles
    # Mount static assets (css, js, fonts) but NOT at root — we'll handle that with a fallback
    for subdir in ["assets", "cursor", "fonts", "scripts"]:
        subpath = os.path.join(FRONTEND_DIR, subdir)
        if os.path.isdir(subpath):
            app.mount(f"/{subdir}", StaticFiles(directory=subpath), name=f"static_{subdir}")
    # Serve materialdesignicons.min.css
    _frontend_css = os.path.join(FRONTEND_DIR, "materialdesignicons.min.css")
else:
    print(f"Frontend build not found at {FRONTEND_DIR} — run 'cd frontend && pnpm build'")
    _frontend_css = None


@app.on_event("startup")
async def load_extensions_on_startup():
    _load_extension_files()

    # Show ready message with runner info
    import httpx
    fal_key = os.environ.get("FAL_KEY", "")
    endpoint = os.environ.get("FAL_COMFY_ENDPOINT", "")
    org = endpoint.split("/")[0] if "/" in endpoint else ""
    app_name = endpoint.split("/")[1] if "/" in endpoint else ""
    dashboard = f"https://fal.ai/dashboard/apps/{org}/{app_name}/runners"

    runner_link = dashboard
    try:
        async with httpx.AsyncClient(timeout=5) as client:
            resp = await client.get(
                "https://rest.alpha.fal.ai/runners/",
                headers={"Authorization": f"Key {fal_key}"},
            )
            runners = resp.json().get("runners", [])
            active = [r for r in runners if r.get("state") in ("IDLE", "ACTIVE", "SETUP")]
            if active:
                rid = active[0]["runner_id"]
                runner_link = f"{dashboard}?s_runner={rid}"
    except Exception:
        pass

    print("")
    print("  ============================================")
    print("  READY — Open http://localhost:8188")
    print(f"  Runner logs: {runner_link}")
    print("  ============================================")
    print("")


class FileStorage(dict):
    _path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "userdata", "file_storage.json")

    @staticmethod
    def get_key(filename: str, subfolder: str, _type: str):
        import unicodedata
        # Normalize Unicode spaces (macOS screenshots use narrow no-break space \u202f)
        filename = unicodedata.normalize("NFKC", filename).replace("\u202f", " ").replace("\u00a0", " ")
        return f"{_type}/{subfolder}/{filename}"

    def _save(self):
        import json
        os.makedirs(os.path.dirname(self._path), exist_ok=True)
        with open(self._path, "w") as f:
            json.dump(dict(self), f)

    def _load(self):
        import json
        if os.path.isfile(self._path):
            with open(self._path, "r") as f:
                self.update(json.load(f))

    def __setitem__(self, key, value):
        super().__setitem__(key, value)
        self._save()


file_storage = FileStorage()
file_storage._load()



class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
        self.client_sessions: Dict[str, WebSocket] = {}

    async def connect(self, websocket: WebSocket, client_id: Optional[str] = None):
        await websocket.accept()
        self.active_connections.append(websocket)

        if not client_id:
            client_id = str(uuid.uuid4())

        self.client_sessions[client_id] = websocket

        # Send initial status message with session ID
        await websocket.send_text(
            json.dumps(
                {
                    "type": "status",
                    "data": {
                        "sid": client_id,
                        "status": {"exec_info": {"queue_remaining": 0}},
                    },
                }
            )
        )

        return client_id

    def disconnect(self, websocket: WebSocket, client_id: Optional[str] = None):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
        if client_id and client_id in self.client_sessions:
            del self.client_sessions[client_id]

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        """Broadcast message to all active connections, handling failures gracefully"""
        if not self.active_connections:
            print("No active WebSocket connections to broadcast to")
            return

        failed_connections = []

        for connection in self.active_connections[:]:  # Create a copy to iterate over
            try:
                await connection.send_text(message)
            except Exception as e:
                print(f"Failed to send message to WebSocket connection: {e}")
                failed_connections.append(connection)

        # Clean up failed connections
        for failed_connection in failed_connections:
            if failed_connection in self.active_connections:
                self.active_connections.remove(failed_connection)
                print("Removed failed WebSocket connection")


manager = ConnectionManager()


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, clientId: Optional[str] = None):
    client_id = await manager.connect(websocket, clientId)

    async def heartbeat():
        """Send a periodic status ping so the browser keepalive doesn't
        classify the connection as dead (silences 'Reconnecté' toasts)."""
        try:
            while True:
                await asyncio.sleep(15)
                await websocket.send_text(
                    json.dumps({
                        "type": "status",
                        "data": {"status": {"exec_info": {"queue_remaining": 0}}},
                    })
                )
        except Exception:
            pass

    hb_task = asyncio.create_task(heartbeat())
    try:
        while True:
            data = await websocket.receive_text()
            print(f"Received from client {client_id}: {data}")

    except WebSocketDisconnect:
        manager.disconnect(websocket, client_id)
        print(f"Client {client_id} disconnected")
    finally:
        hb_task.cancel()


# Models
class ComfyApiWorkflow(BaseModel):
    pass


class ComfyWorkflowJSON(dict):
    pass


@dataclass
class FalComfyInput:
    prompt: dict = field(default_factory=dict)
    fal_inputs: dict = field(default_factory=dict)
    fal_inputs_dev_info: dict = field(default_factory=dict)


@dataclass
class FalComfyOutput:
    outputs: dict = field(default_factory=dict)
    fal_outputs: dict = field(default_factory=dict)
    timings: dict = field(default_factory=dict)


class QueuePromptRequestBody(BaseModel):
    client_id: str | None = None
    prompt: dict
    extra_data: Dict[str, Any] | None = None
    front: Optional[bool] = None
    number: Optional[int] = None


@dataclass
class SystemStats:
    python_version: str
    os: str
    devices: List[Dict[str, Any]] = Field(default_factory=list)


class Settings(BaseModel):
    pass


class User(BaseModel):
    username: str


# In-memory storage
queue_running = []
queue_pending = []
history = []
client_sessions = {}
websocket_connections = set()
# Add task tracking for running fal_client_task instances
running_tasks: Dict[str, asyncio.Task] = {}

# Output storage — persists generated file URLs
class OutputStorage:
    _path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "userdata", "outputs.json")
    _outputs: list = []

    @classmethod
    def load(cls):
        import json
        if os.path.isfile(cls._path):
            with open(cls._path, "r") as f:
                cls._outputs = json.load(f)

    @classmethod
    def save(cls):
        import json
        os.makedirs(os.path.dirname(cls._path), exist_ok=True)
        with open(cls._path, "w") as f:
            json.dump(cls._outputs, f)

    @classmethod
    def add(cls, prompt_id: str, urls: list, metadata: dict = None):
        from datetime import datetime
        # Skip URLs that already exist anywhere in storage to avoid duplicate
        # cards in the UI when an `executed` event fires twice for the same
        # output (re-queue, runtime retry, etc.).
        existing = {u for e in cls._outputs for u in e.get("urls", [])}
        new_urls = [u for u in urls if u not in existing]
        if not new_urls:
            return
        entry = {
            "prompt_id": prompt_id,
            "urls": new_urls,
            "timestamp": datetime.utcnow().isoformat(),
            "metadata": metadata or {},
        }
        cls._outputs.append(entry)
        cls.save()

    @classmethod
    def get_all(cls):
        return cls._outputs

OutputStorage.load()


# --- Extension JS files cache ---
_extension_files_cache: Dict[str, str] = {}
_extensions_list_cache: list = []


def _load_extension_files():
    """Download all extension JS files from fal server and cache locally."""
    global _extension_files_cache, _extensions_list_cache
    import httpx
    try:
        print("Fetching extension JS files from server...")
        resp = httpx.post(
            f"https://fal.run/{ENDPOINT}/extension_files",
            headers={"Authorization": f"Key {FAL_KEY}"},
            json={},
            timeout=120.0,
        )
        resp.raise_for_status()
        data = resp.json()
        if isinstance(data, dict) and not data.get("detail"):
            _extension_files_cache = data
            _extensions_list_cache = list(data.keys())
            print(f"Cached {len(data)} extension JS files")
        else:
            print(f"Unexpected extension_files response: {data.get('detail', str(data)[:200])}")
    except Exception as e:
        print(f"Failed to load extension files: {e}")
    # Fallback: get just the list
    if not _extension_files_cache:
        try:
            resp = httpx.post(
                f"https://fal.run/{ENDPOINT}/extensions",
                headers={"Authorization": f"Key {FAL_KEY}"},
                json={},
                timeout=30.0,
            )
            resp.raise_for_status()
            ext_list = resp.json()
            if isinstance(ext_list, list):
                _extensions_list_cache = ext_list
                print(f"Got extension list ({len(_extensions_list_cache)} entries) but not file contents")
        except Exception as e2:
            print(f"Failed to load extension list: {e2}")


FAL_EXTRA_JS = "/extensions/fal-extras/extras.js"


@app.get("/api/extensions")
async def get_extensions():
    if not _extensions_list_cache:
        _load_extension_files()
    out = list(_extensions_list_cache)
    if FAL_EXTRA_JS not in out:
        out.append(FAL_EXTRA_JS)
    return out


@app.get(FAL_EXTRA_JS)
async def get_fal_extras_js():
    """Serve the local fal extras JS (inpaint dropdown + right-click save image)."""
    p = Path(__file__).resolve().parent.parent / "frontend" / "fal-extras" / "extras.js"
    if not p.exists():
        raise HTTPException(404, "fal-extras/extras.js not found")
    return Response(content=p.read_text(), media_type="application/javascript")


@app.get("/extensions/{path:path}")
async def get_extension_file(path: str):
    if not _extension_files_cache:
        _load_extension_files()

    # Try multiple key formats + common aliases
    candidates = [
        f"/extensions/{path}",
        f"extensions/{path}",
        path,
    ]
    # Handle rgthree -> rgthree-comfy alias (files are under rgthree-comfy/comfyui/)
    if path.startswith("rgthree/"):
        rest = path[len("rgthree/"):]
        for prefix in ["rgthree-comfy/comfyui/", "rgthree-comfy/", "rgthree-comfy/common/"]:
            alt = prefix + rest
            candidates.extend([f"/extensions/{alt}", f"extensions/{alt}", alt])
    # Handle scripts/xxx.js -> serve from frontend build
    if path.startswith("scripts/"):
        script_path = os.path.join(FRONTEND_DIR, path)
        if os.path.isfile(script_path):
            return Response(
                content=open(script_path, "rb").read(),
                media_type="application/javascript",
            )
    # Handle core/xxx.js -> serve from frontend build
    if path.startswith("core/"):
        core_path = os.path.join(FRONTEND_DIR, "extensions", path)
        if os.path.isfile(core_path):
            return Response(
                content=open(core_path, "rb").read(),
                media_type="application/javascript",
            )

    for key in candidates:
        content = _extension_files_cache.get(key)
        if content is not None:
            return Response(
                content=content,
                media_type="application/javascript",
            )

    raise HTTPException(404, detail=f"Extension file not found: {path}")


@app.get("/api/workflow_templates")
async def get_workflow_templates():
    return {}


@app.get("/templates/index.json")
async def get_core_workflow_templates():
    return []


@app.get("/api/embeddings")
async def get_embeddings():
    return {"embeddings": []}


_object_info_cache = {"data": None, "gzipped": None}


@app.get("/api/object_info")
async def get_node_defs() -> Response:
    import gzip

    # Pre-compute and cache gzipped bytes (node defs don't change during a session)
    if _object_info_cache["gzipped"] is None:
        json_data = json.dumps(NODE_CLASS_MAPPINGS).encode("utf-8")
        _object_info_cache["gzipped"] = gzip.compress(json_data, compresslevel=6)

    compressed_data = _object_info_cache["gzipped"]

    return Response(
        content=compressed_data,
        media_type="application/json",
        headers={
            "Content-Encoding": "gzip",
            "Content-Length": str(len(compressed_data)),
            "Cache-Control": "private, max-age=60",
        },
    )


@app.get("/api/object_info/{node_class}")
async def get_node_def(node_class: str) -> dict:
    return NODE_CLASS_MAPPINGS.get(node_class, {})


@app.post("/api/prompt")
async def queue_prompt(request: QueuePromptRequestBody) -> JSONResponse:
    # Generate a prompt ID for this execution
    prompt_id = str(uuid.uuid4())

    try:
        # Track this prompt in queue_running
        queue_running.append([len(queue_running) + len(queue_pending) + 1, prompt_id, request.prompt, {"client_id": request.client_id}, []])

        # Send prompt queued event immediately
        await manager.broadcast(
            json.dumps(
                {
                    "type": "promptQueued",
                    "data": {"number": len(queue_running), "batchCount": 1},
                }
            )
        )

        # Replace plain filenames with fal CDN URLs for load nodes
        # and build fal_inputs for proper pre-download handling
        prompt = request.prompt
        fal_inputs = {}
        fal_inputs_dev_info = {}
        load_nodes = {
            "LoadImage": "image",
            "VHS_LoadVideo": "video",
            "VHS_LoadVideoPath": "video",
            "VHS_LoadAudio": "audio",
            "VHS_LoadAudioUpload": "audio",
            "LoadAudio": "audio",
        }
        for node_id, node_data in prompt.items():
            class_type = node_data.get("class_type", "")
            input_name = load_nodes.get(class_type)
            if input_name:
                raw_val = node_data.get("inputs", {}).get(input_name, "")
                if raw_val and not str(raw_val).startswith(("https://", "http://", "data:")):
                    # ComfyUI annotates filenames as "<path> [<type>]" e.g.
                    # "clipspace/image.png [input]". Strip that annotation before
                    # the file_storage lookup, otherwise the key never matches.
                    val = re.sub(r"\s*\[(input|output|temp)\]\s*$", "", str(raw_val))
                    # Look up the fal CDN URL from file_storage
                    for _type in ("input", "output", "temp"):
                        for subfolder in ("", val.rsplit("/", 1)[0] if "/" in val else ""):
                            fname = val.rsplit("/", 1)[-1] if "/" in val else val
                            key = FileStorage.get_key(fname, subfolder, _type)
                            if key in file_storage:
                                field = str(uuid.uuid4())
                                fal_inputs[field] = file_storage[key]
                                fal_inputs_dev_info[field] = {
                                    "key": [node_id, "inputs", input_name],
                                    "class_type": class_type,
                                }
                                break

        # Start the fal client connection to check status immediately
        fal_workflow = FalComfyInput(prompt=prompt)
        if fal_inputs:
            fal_workflow.fal_inputs = fal_inputs
            fal_workflow.fal_inputs_dev_info = fal_inputs_dev_info
            print(f"[proxy] fal_inputs: {fal_inputs}")
            print(f"[proxy] fal_inputs_dev_info: {fal_inputs_dev_info}")

        payload = asdict(fal_workflow)
        print(f"[proxy] Sending to fal, keys: {list(payload.keys())}, has fal_inputs: {'fal_inputs' in payload and bool(payload['fal_inputs'])}")

        # Connect to fal — NO retry loop to avoid spawning multiple runners.
        # If the runner is not ready, we return an error to the frontend and let
        # the user click Run again. Each POST can spawn a new runner, so retrying
        # here would leak runners.
        event_source_ctx = None
        try:
            event_source = aconnect_sse(
                fal_client.async_client._client,
                "POST",
                f"{RUN_URL_FORMAT}{ENDPOINT}/stream",
                json=payload,
            )
            event_source_ctx = await event_source.__aenter__()
        except Exception as e:
            return JSONResponse(
                status_code=400,
                content={
                    "error": f"Failed to connect to fal: {str(e)}",
                    "prompt_id": prompt_id,
                    "node_errors": {},
                },
            )

        try:
            if event_source_ctx.response.status_code != 200:
                # Handle error immediately and return to client
                try:
                    error_message = await event_source_ctx.response.aread()
                    error_json = json.loads(error_message)
                    error_details = error_json.get("detail", [])

                    # Clean up the connection
                    await event_source.__aexit__(None, None, None)

                    return JSONResponse(status_code=400, content=error_details)
                except Exception as parse_error:
                    import traceback

                    traceback.print_exc()
                    print(parse_error)
                    print(f"Failed to parse error response: {parse_error}")

                    # Clean up the connection
                    await event_source.__aexit__(None, None, None)

                    return JSONResponse(
                        status_code=400,
                        content={
                            "error": "Failed to process workflow",
                            "prompt_id": prompt_id,
                            "node_errors": {},
                        },
                    )

            # Start background task to continue processing remaining events
            task = asyncio.create_task(
                fal_client_task_with_source(prompt_id, event_source, event_source_ctx)
            )
            # Store the task reference for potential cancellation
            running_tasks[prompt_id] = task

            print(f"Successfully queued prompt with ID: {prompt_id}")

            return JSONResponse(
                {
                    "prompt_id": prompt_id,
                    "number": len(queue_pending) + 1,
                    "node_errors": {},
                }
            )

        except Exception as e:
            # Clean up the connection on any error
            try:
                await event_source.__aexit__(None, None, None)
            except:
                pass
            raise e

    except Exception as e:
        import traceback
        print(f"Failed to queue prompt {prompt_id}: {str(e)}")
        traceback.print_exc()

        # Clean up if task was created but failed
        if prompt_id in running_tasks:
            running_tasks[prompt_id].cancel()
            del running_tasks[prompt_id]

        # Remove from queue_running on failure
        for i, item in enumerate(queue_running):
            if len(item) > 1 and item[1] == prompt_id:
                queue_running.pop(i)
                break

        raise HTTPException(status_code=500, detail=f"Failed to queue prompt: {str(e)}")


async def fal_client_task_with_source(prompt_id: str, event_source, event_source_ctx):
    """Continue processing an existing fal_client event source with proper error handling"""
    try:
        print(f"Starting fal_client_task_with_source for prompt_id: {prompt_id}")

        fal_request_id = event_source_ctx.response.headers.get("x-fal-request-id")
        runner_id_hdr = event_source_ctx.response.headers.get("x-fal-runner-id")
        if runner_id_hdr:
            # Record the runner currently serving this proxy (used by admin UI)
            globals()["MY_RUNNER_ID"] = runner_id_hdr
        print(f"fal_request_id: {fal_request_id}")

        # Process events from the already-opened event source
        async for event_sse in event_source_ctx.aiter_sse():
            event = event_sse.json()

            # Check if task was cancelled
            current_task = asyncio.current_task()
            if current_task and current_task.cancelled():
                print(f"Task cancelled for prompt_id: {prompt_id}")
                break

            try:
                # Capture output URLs — only "output" type (skip previews & temp files)
                if event.get("type") == "executed":
                    output_data = event.get("data", {}).get("output", {})
                    urls = []
                    for key in ("images", "gifs"):
                        for item in output_data.get(key, []):
                            # Skip previews and temp files
                            if item.get("type") != "output":
                                continue
                            if "url" in item:
                                urls.append(item["url"])
                    if urls:
                        OutputStorage.add(prompt_id, urls, {"node": event.get("data", {}).get("node")})

                if event["type"] == "fal-execution-error":
                    # Format as ComfyUI execution_error so the frontend highlights the failing node
                    err = event.get("data", {}).get("error", {}).get("error", {})
                    node_id = err.get("node_id") or event.get("data", {}).get("node_id") or err.get("details", {}).get("node_id")
                    node_type = err.get("node_type") or err.get("details", {}).get("node_type")
                    msg = err.get("exception_message") or err.get("message") or str(err)
                    ex_type = err.get("exception_type") or "ExecutionError"
                    traceback_info = err.get("traceback") or []
                    formatted = {
                        "type": "execution_error",
                        "data": {
                            "prompt_id": prompt_id,
                            "node_id": node_id,
                            "node_type": node_type,
                            "exception_message": msg,
                            "exception_type": ex_type,
                            "traceback": traceback_info,
                            "current_inputs": err.get("current_inputs", {}),
                            "current_outputs": err.get("current_outputs", []),
                        },
                    }
                    await manager.broadcast(json.dumps(formatted))
                    print(f"[proxy] Execution error on node {node_id} ({node_type}): {msg}")
                else:
                    await manager.broadcast(json.dumps(event))
            except Exception as broadcast_error:
                print(
                    f"Failed to broadcast event for prompt_id {prompt_id}: {broadcast_error}"
                )
                # Continue processing even if broadcast fails

    except asyncio.CancelledError:
        print(f"Task was cancelled for prompt_id: {prompt_id}")
        # Task was cancelled, send cancellation event
        try:
            await manager.broadcast(
                json.dumps(
                    {"type": "execution_interrupted", "data": {"prompt_id": prompt_id}}
                )
            )
        except Exception as e:
            print(f"Failed to broadcast cancellation for prompt_id {prompt_id}: {e}")
        raise  # Re-raise to maintain cancellation semantics

    except Exception as e:
        import traceback
        print(
            f"Error in fal_client_task_with_source for prompt_id {prompt_id}: {str(e)}",
        )

        # Handle other errors - broadcast error notification
        error_message = {
            "type": "execution_error",
            "data": {
                "prompt_id": prompt_id,
                "error": str(e),
                "node_id": None,
                "node_type": None,
            },
        }

        try:
            await manager.broadcast(json.dumps(error_message))
            print(f"Successfully broadcasted error for prompt_id: {prompt_id}")
        except Exception as broadcast_error:
            print(
                f"Failed to broadcast error for prompt_id {prompt_id}: {broadcast_error}"
            )

        # IMPORTANT: Re-raise the exception so the task is marked as failed
        raise

    finally:
        # Clean up the event source connection
        try:
            await event_source.__aexit__(None, None, None)
        except Exception as cleanup_error:
            print(
                f"Failed to cleanup event source for prompt_id {prompt_id}: {cleanup_error}"
            )

        # Clean up the task reference
        if prompt_id in running_tasks:
            del running_tasks[prompt_id]
            print(f"Cleaned up task reference for prompt_id: {prompt_id}")

        # Remove from queue_running
        for i, item in enumerate(queue_running):
            if len(item) > 1 and item[1] == prompt_id:
                queue_running.pop(i)
                break


@app.get("/api/vhs/viewvideo")
async def vhs_viewvideo(request: Request):
    """VHS advanced video preview with local ffmpeg processing."""
    import httpx, subprocess, hashlib
    params = dict(request.query_params)
    filename = params.get("filename", "")
    subfolder = params.get("subfolder", "")
    file_type = params.get("type", "input")

    print(f"[VHS viewvideo] filename={filename} params={params}")

    # Find the fal.media URL
    key = FileStorage.get_key(filename, subfolder, file_type)
    url = file_storage.get(key)
    if not url:
        url = params.get("url")
    if not url:
        print(f"[VHS viewvideo] ERROR: video not found in storage. key={key}")
        print(f"[VHS viewvideo] Available keys: {list(file_storage.keys())}")
        raise HTTPException(400, "Video not found")

    # Cache the source video locally
    cache_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "userdata", "video_cache")
    os.makedirs(cache_dir, exist_ok=True)
    url_hash = hashlib.md5(url.encode()).hexdigest()
    ext = filename.rsplit(".", 1)[-1] if "." in filename else "mp4"
    cached_path = os.path.join(cache_dir, f"{url_hash}.{ext}")

    if not os.path.isfile(cached_path):
        print(f"[VHS viewvideo] Downloading {url} -> {cached_path}")
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            with open(cached_path, "wb") as f:
                f.write(resp.content)
        print(f"[VHS viewvideo] Downloaded {os.path.getsize(cached_path)} bytes")
    else:
        print(f"[VHS viewvideo] Using cached {cached_path}")

    # Parse params
    skip_first_frames = int(float(params.get("skip_first_frames", 0)))
    frame_load_cap = int(float(params.get("frame_load_cap", 0)))
    force_rate = float(params.get("force_rate", 0))
    force_size = params.get("force_size", "Disabled")
    select_every_nth = int(float(params.get("select_every_nth", 1)))
    deadline = params.get("deadline", "realtime")

    # Get source video fps via ffprobe
    try:
        probe = subprocess.run(
            ["ffprobe", "-v", "error", "-select_streams", "v:0",
             "-show_entries", "stream=r_frame_rate", "-of", "csv=p=0", cached_path],
            capture_output=True, text=True, timeout=10
        )
        num, den = probe.stdout.strip().split("/")
        source_fps = float(num) / float(den)
    except Exception:
        source_fps = 30.0

    target_rate = force_rate if force_rate > 0 else source_fps
    print(f"[VHS viewvideo] source_fps={source_fps} target_rate={target_rate} skip={skip_first_frames} cap={frame_load_cap} force_size={force_size}")

    # Build ffmpeg args
    cmd = ["ffmpeg", "-v", "error", "-y"]

    # Seek to skip_first_frames
    if skip_first_frames > 0:
        start_time = skip_first_frames / target_rate
        cmd += ["-ss", f"{start_time:.3f}"]

    cmd += ["-i", cached_path]

    if force_rate > 0:
        cmd += ["-r", str(force_rate)]

    # Scale filter — parse force_size like "510x?" or "510x380"
    filters = []
    if force_size and force_size != "Disabled":
        parts = force_size.split("x")
        if len(parts) == 2:
            w = parts[0].strip()
            h = parts[1].strip()
            if h == "?":
                w_int = int(float(w))
                # Ensure even
                w_int = w_int + (w_int % 2)
                filters.append(f"scale={w_int}:-2")
            elif w == "?":
                h_int = int(float(h))
                h_int = h_int + (h_int % 2)
                filters.append(f"scale=-2:{h_int}")
            else:
                try:
                    w_int = int(float(w))
                    h_int = int(float(h))
                    # Ensure even dimensions
                    w_int = w_int + (w_int % 2)
                    h_int = h_int + (h_int % 2)
                    filters.append(f"scale={w_int}:{h_int}")
                except ValueError:
                    pass

    if select_every_nth > 1:
        filters.append(f"select=not(mod(n\\,{select_every_nth}))")

    if filters:
        cmd += ["-vf", ",".join(filters)]

    if frame_load_cap > 0:
        cmd += ["-frames:v", str(frame_load_cap)]

    cmd += ["-c:v", "libvpx-vp9", "-deadline", deadline, "-cpu-used", "8",
            "-f", "webm", "pipe:1"]

    print(f"[VHS viewvideo] Running: {' '.join(cmd)}")

    try:
        result = subprocess.run(cmd, capture_output=True, timeout=60)
        if result.returncode != 0:
            print(f"[VHS viewvideo] ffmpeg FAILED: {result.stderr.decode()}")
            return await view_image(request)
        print(f"[VHS viewvideo] OK: {len(result.stdout)} bytes output")
        return Response(content=result.stdout, media_type="video/webm")
    except Exception as e:
        print(f"[VHS viewvideo] Exception: {e}")
        return await view_image(request)


@app.get("/api/experiment/models")
async def get_model_folders() -> list[dict[str, str | list[str]]]:
    model_folders = [{"name": k, "folders": [k]} for k in MODELS.keys()]
    return model_folders


@app.get("/api/experiment/models/{folder}")
async def get_models(folder: str) -> list[dict[str, str | int]]:
    return MODELS.get(folder, [])


@app.get("/api/view_metadata/{folder}")
async def view_metadata(folder: str, filename: str):
    return None


@app.get("/api/view")
async def view_image(request: Request):
    """Simple view function that handles URL-based file serving from fal.ai"""
    params = dict(request.query_params)
    filename = params.get("filename", "")
    subfolder = params.get("subfolder", "")
    file_type = params.get("type", "output")
    url = params.get("url")
    preview = params.get("preview")
    channel = params.get("channel")

    # ComfyUI annotates filenames as "<name> [<type>]" (e.g. "image.png [input]")
    # Strip the annotation before looking up file_storage.
    m = re.match(r"^(.*?)\s*\[(input|output|temp)\]\s*$", filename)
    if m:
        filename = m.group(1)
        file_type = m.group(2)

    # Resolve URL: either from query param or from file_storage
    if not url:
        key = FileStorage.get_key(filename, subfolder, file_type)
        if key in file_storage:
            url = file_storage[key]

    if not url:
        raise HTTPException(status_code=400, detail="File not found")

    # For videos/media: proxy the bytes directly (no Image.open)
    VIDEO_EXTS = (".mp4", ".gif", ".webm", ".mov", ".avi", ".mkv")
    if any(filename.lower().endswith(ext) for ext in VIDEO_EXTS) or any(url.lower().endswith(ext) for ext in VIDEO_EXTS):
        import httpx
        async with httpx.AsyncClient(timeout=180) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            ext = filename.rsplit(".", 1)[-1].lower()
            mime = {"mp4": "video/mp4", "gif": "image/gif", "webm": "video/webm",
                    "mov": "video/quicktime", "avi": "video/x-msvideo", "mkv": "video/x-matroska"
                    }.get(ext, "application/octet-stream")
            return Response(content=resp.content, media_type=mime,
                            headers={"Content-Disposition": f'filename="{filename}"'})

    return await handle_url_image(url, filename, preview, channel)


async def handle_url_image(
    url: str,
    filename: str,
    preview: Optional[str] = None,
    channel: Optional[str] = None,
):
    """Handle image from URL with optional processing"""
    try:
        import httpx

        # Fetch the image from the URL
        async with httpx.AsyncClient() as client:
            response = await client.get(url)
            response.raise_for_status()
            image_data = response.content

        # Save to S3 in the background (non-blocking)
        if s3_client:
            content_type = response.headers.get("content-type", "image/png")
            asyncio.get_event_loop().run_in_executor(
                None, lambda: s3_upload_sync(image_data, filename, content_type, "outputs")
            )

        # If no processing needed, return the original image
        if not preview and not channel:
            content_type = response.headers.get("content-type", "image/png")
            return Response(
                content=image_data,
                media_type=content_type,
                headers={"Content-Disposition": f'filename="{filename}"'},
            )

        # For processing, load the image
        image_buffer = BytesIO(image_data)
        img = Image.open(image_buffer)

        # Handle preview mode
        if preview:
            return await process_preview(img, filename, preview, channel)

        # Handle channel extraction
        if channel:
            return await process_channel(img, filename, channel)

        # Default: return original
        content_type = response.headers.get("content-type", "image/png")
        return Response(
            content=image_data,
            media_type=content_type,
            headers={"Content-Disposition": f'filename="{filename}"'},
        )

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Failed to process image from URL: {str(e)}"
        )


async def process_preview(
    img: Image.Image, filename: str, preview: str, channel: Optional[str]
):
    """Process image preview with format and quality options"""
    try:
        preview_info = preview.split(";")
        image_format = preview_info[0]

        # Validate format
        if image_format not in ["webp", "jpeg"] or (channel and "a" in channel):
            image_format = "webp"

        # Extract quality
        quality = 90
        if len(preview_info) > 1 and preview_info[-1].isdigit():
            quality = int(preview_info[-1])

        # Convert image if needed
        if image_format == "jpeg" or (channel and channel == "rgb"):
            if img.mode in ("RGBA", "LA", "P"):
                img = img.convert("RGB")

        # Save processed image
        buffer = BytesIO()
        img.save(buffer, format=image_format.upper(), quality=quality)
        buffer.seek(0)

        return Response(
            content=buffer.getvalue(),
            media_type=f"image/{image_format}",
            headers={"Content-Disposition": f'filename="{filename}"'},
        )

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Preview processing failed: {str(e)}"
        )


async def process_channel(img: Image.Image, filename: str, channel: str):
    """Process channel extraction"""
    try:
        buffer = BytesIO()

        if channel == "rgb":
            # Extract RGB channels
            if img.mode == "RGBA":
                r, g, b, a = img.split()
                rgb_img = Image.merge("RGB", (r, g, b))
            else:
                rgb_img = img.convert("RGB")
            rgb_img.save(buffer, format="PNG")

        elif channel == "a":
            # Extract alpha channel
            if img.mode == "RGBA":
                r, g, b, a = img.split()
                alpha_img = Image.new("RGBA", img.size)
                alpha_img.putalpha(a)
            else:
                # No alpha channel, create a white alpha
                alpha_img = Image.new("RGBA", img.size, (255, 255, 255, 255))
            alpha_img.save(buffer, format="PNG")

        else:
            # Default to full image
            img.save(buffer, format="PNG")

        buffer.seek(0)

        return Response(
            content=buffer.getvalue(),
            media_type="image/png",
            headers={"Content-Disposition": f'filename="{filename}"'},
        )

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Channel extraction failed: {str(e)}"
        )


@app.get("/api/my_runner")
async def my_runner():
    """Return the runner ID currently serving this proxy (captured from x-fal-runner-id)."""
    return {"runner_id": globals().get("MY_RUNNER_ID")}


@app.get("/api/outputs")
async def get_outputs():
    """JSON list of all generated outputs (URLs are deduplicated across runs)."""
    seen: set[str] = set()
    deduped: list[dict] = []
    for entry in OutputStorage.get_all():
        unique_urls: list[str] = []
        for u in entry.get("urls", []):
            if u in seen:
                continue
            seen.add(u)
            unique_urls.append(u)
        if unique_urls:
            deduped.append({**entry, "urls": unique_urls})
    return deduped


@app.get("/outputs")
async def outputs_page():
    """HTML page to browse and download all generated outputs."""
    outputs = OutputStorage.get_all()
    rows = ""
    for entry in reversed(outputs):
        for url in entry.get("urls", []):
            fname = url.rsplit("/", 1)[-1] if "/" in url else "output"
            ext = fname.rsplit(".", 1)[-1] if "." in fname else ""
            is_video = ext in ("mp4", "gif", "webm", "mov")
            preview = f'<video src="{url}" controls muted loop style="max-width:300px;max-height:200px"></video>' if is_video else f'<img src="{url}" style="max-width:300px;max-height:200px">'
            dl_url = f"/api/download?url={url}&filename={fname}"
            rows += f'<tr><td>{preview}</td><td>{entry["timestamp"][:19]}</td><td>{ext.upper()}</td><td><a href="{dl_url}">Download</a> · <a href="{url}" target="_blank">Open</a></td></tr>\n'
    total_files = sum(len(e.get("urls", [])) for e in outputs)
    html = f"""<!DOCTYPE html>
<html><head><title>ComfyUI Outputs</title>
<style>body{{font-family:sans-serif;padding:20px;background:#1a1a1a;color:#eee}}
table{{border-collapse:collapse;width:100%}}td,th{{padding:10px;border-bottom:1px solid #333;text-align:left}}
a{{color:#4fc3f7;margin-right:8px}}h1{{color:#fff}}h2{{color:#aaa;font-weight:normal;font-size:14px}}</style></head>
<body><h1>ComfyUI Outputs</h1><h2>{total_files} files from {len(outputs)} runs</h2>
<table><tr><th>Preview</th><th>Time</th><th>Type</th><th>Actions</th></tr>
{rows if rows else '<tr><td colspan="4">No outputs yet. Run a workflow first.</td></tr>'}
</table></body></html>"""
    return HTMLResponse(content=html)


@app.get("/api/download")
async def download_file(url: str, filename: str = "output"):
    """Proxy download from fal.media to force browser download."""
    import httpx
    async with httpx.AsyncClient(timeout=180) as client:
        resp = await client.get(url)
        resp.raise_for_status()
        content_type = resp.headers.get("content-type", "application/octet-stream")
        return Response(
            content=resp.content,
            media_type=content_type,
            headers={"Content-Disposition": f'attachment; filename="{filename}"'},
        )


@app.get("/api/queue")
async def get_queue():
    return {"queue_running": queue_running, "queue_pending": queue_pending}


@app.get("/api/history")
async def get_history(max_items: int = 200):
    return history[:max_items]


@app.get("/api/system_stats")
async def get_system_stats():
    return SystemStats(
        python_version="3.9.0", os="Linux", devices=[{"name": "CPU", "type": "cpu"}]
    )


@app.post("/api/interrupt")
async def interrupt():
    """Cancel all running fal_client_task instances and notify the remote server"""
    cancelled_count = 0

    # Cancel all local running tasks
    for prompt_id, task in list(running_tasks.items()):
        if not task.done():
            task.cancel()
            cancelled_count += 1

    # Clear the running tasks dictionary
    running_tasks.clear()

    # Also cancel on the remote fal server
    try:
        import httpx
        async with httpx.AsyncClient() as client:
            await client.post(
                f"{RUN_URL_FORMAT}{ENDPOINT}/cancel",
                headers={"Authorization": f"Key {FAL_KEY}"},
                timeout=5.0,
            )
    except Exception as e:
        print(f"Failed to cancel remote execution: {e}")

    # Broadcast interrupt event to all connected clients
    await manager.broadcast(
        json.dumps(
            {
                "type": "status",
                "data": {"status": {"exec_info": {"queue_remaining": 0}}},
            }
        )
    )

    return {"status": "success", "cancelled_tasks": cancelled_count}


@app.get("/api/users")
async def get_user_config():
    return {"username": "default_user"}


@app.post("/api/users")
async def create_user(user: User):
    return {"username": user.username}


@app.get("/api/settings")
async def get_settings():
    return {}


@app.get("/api/settings/{setting_id}")
async def get_setting(setting_id: str):
    return None


@app.post("/api/settings")
async def store_settings(settings: Settings):
    return {"status": "success"}


@app.post("/api/settings/{setting_id}")
async def store_setting(setting_id: str, request: Request):
    return {"status": "success"}


# Local userdata storage directory
USERDATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "userdata")
os.makedirs(USERDATA_DIR, exist_ok=True)


@app.get("/api/userdata/{file:path}")
async def get_user_data(file: str):
    file_path = os.path.join(USERDATA_DIR, file)
    if os.path.isfile(file_path):
        return Response(
            content=open(file_path, "rb").read(),
            media_type="application/json",
        )
    return {}


@app.post("/api/userdata/{file:path}")
async def store_user_data(file: str, request: Request):
    file_path = os.path.join(USERDATA_DIR, file)
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    body = await request.body()
    with open(file_path, "wb") as f:
        f.write(body)
    return {"status": "success"}


@app.delete("/api/userdata/{file:path}")
async def delete_user_data(file: str):
    file_path = os.path.join(USERDATA_DIR, file)
    if os.path.isfile(file_path):
        os.remove(file_path)
    return {"status": "success"}


@app.post("/api/userdata/{source:path}/move/{dest:path}")
async def move_user_data(source: str, dest: str, overwrite: bool = False):
    src_path = os.path.join(USERDATA_DIR, source)
    dst_path = os.path.join(USERDATA_DIR, dest)
    if os.path.isfile(src_path):
        os.makedirs(os.path.dirname(dst_path), exist_ok=True)
        os.rename(src_path, dst_path)
    return {"status": "success"}


@app.get("/api/userdata")
async def list_user_data_full_info(
    dir: str = "", recurse: bool = True, split: bool = False, full_info: bool = True
):
    target = os.path.join(USERDATA_DIR, dir)
    if not os.path.isdir(target):
        return []
    files = []
    for root, dirs, filenames in os.walk(target):
        for fname in filenames:
            full = os.path.join(root, fname)
            rel = os.path.relpath(full, USERDATA_DIR)
            if full_info:
                files.append({"path": rel, "size": os.path.getsize(full), "modified": os.path.getmtime(full)})
            else:
                files.append(rel)
        if not recurse:
            break
    return files


@app.get("/internal/logs")
async def get_logs():
    return ""


@app.get("/internal/logs/raw")
async def get_raw_logs():
    return {"logs": []}


@app.patch("/internal/logs/subscribe")
async def subscribe_logs(enabled: bool, clientId: str):
    return {"status": "success"}


@app.get("/internal/folder_paths")
async def get_folder_paths():
    return {}


@app.get("/api/i18n")
async def get_custom_nodes_i18n():
    return {}


@app.get("/api/global_subgraphs")
async def get_global_subgraphs():
    return []


# --- Jobs endpoints (used by the frontend queue store) ---

@app.get("/api/jobs")
async def list_jobs(status: str = "", limit: int = 64, offset: int = 0):
    return {"jobs": [], "pagination": {"total": 0, "offset": offset, "limit": limit, "has_more": False}}


@app.get("/api/jobs/{job_id}")
async def get_job(job_id: str):
    return {"job_id": job_id, "status": "completed"}


@app.post("/api/upload/mask")
async def upload_mask(
    image: UploadFile = File(...),
    original_ref: str = Form(""),
    subfolder: str = Form("clipspace"),
    file_type: str = Form("input", alias="type"),
):
    """Combine the mask painted in ComfyUI's MaskEditor with the original image
    and save the RGBA result so LoadImage's mask output is populated.

    ComfyUI sends multipart with:
      - `image`: the painted mask PNG (RGBA — alpha holds the painted mask)
      - `original_ref`: JSON like {"filename": "...", "type": "...", "subfolder": "..."}

    We rebuild the same RGBA file ComfyUI core produces locally, upload to
    fal, and register it under the original filename so the LoadImage widget
    picks it up.
    """
    import json as _json
    import io
    import unicodedata

    try:
        ref = _json.loads(original_ref) if original_ref else {}
    except Exception:
        ref = {}
    orig_filename = ref.get("filename") or image.filename or "image.png"
    # Fall back to the multipart form fields if `original_ref` doesn't carry
    # subfolder/type — ComfyUI's MaskEditor sends them as form fields.
    orig_subfolder = ref.get("subfolder", subfolder or "")
    orig_type = ref.get("type", file_type or "input")

    # Locate the original image (must already be in file_storage from a prior upload)
    key = FileStorage.get_key(orig_filename, orig_subfolder, orig_type)
    orig_url = file_storage.get(key)
    mask_bytes = await image.read()

    if orig_url is None:
        # Original wasn't found — fall back to storing the mask alone under
        # the original (filename, subfolder, type) so the LoadImage widget
        # value stays the same and the rewriter resolves it.
        url = await fal_client.upload_async(mask_bytes, content_type="image/png")
        out_key = FileStorage.get_key(orig_filename, orig_subfolder, "input")
        file_storage[out_key] = url
        # Also register under "clipspace" subfolder so a workflow that has the
        # ComfyUI-canonical "clipspace/<file>" path still resolves correctly.
        file_storage[FileStorage.get_key(orig_filename, "clipspace", "input")] = url
        return {"name": orig_filename, "subfolder": orig_subfolder, "type": "input", "url": url}

    # Fetch original + composite (RGB from original, A from mask)
    import httpx
    async with httpx.AsyncClient(timeout=60) as client:
        r = await client.get(orig_url)
        r.raise_for_status()
        orig_bytes = r.content

    try:
        orig = Image.open(io.BytesIO(orig_bytes)).convert("RGBA")
        mask_img = Image.open(io.BytesIO(mask_bytes)).convert("RGBA")
        if mask_img.size != orig.size:
            mask_img = mask_img.resize(orig.size, Image.LANCZOS)
        # ComfyUI convention: alpha channel of the saved file holds the mask
        # (white = keep, black = masked area). The MaskEditor sends RGBA with
        # alpha already encoding the painted area.
        r_, g_, b_, _ = orig.split()
        _, _, _, a_ = mask_img.split()
        merged = Image.merge("RGBA", (r_, g_, b_, a_))
        out_buf = io.BytesIO()
        merged.save(out_buf, format="PNG")
        out_buf.seek(0)
        merged_bytes = out_buf.getvalue()
    except Exception as e:
        # If anything fails, keep the mask alone — at least don't break the UI
        merged_bytes = mask_bytes

    # Upload composite, overwrite the original file_storage entry so the
    # workflow's existing LoadImage value (e.g. "image.png") still resolves.
    # Also register the same URL under "clipspace/<file>" so workflows whose
    # widget value was rewritten to ComfyUI's canonical clipspace path resolve
    # too. Cover both shapes — let the rewriter pick whichever the prompt has.
    url = await fal_client.upload_async(merged_bytes, content_type="image/png")
    file_storage[FileStorage.get_key(orig_filename, orig_subfolder, "input")] = url
    file_storage[FileStorage.get_key(orig_filename, "clipspace", "input")] = url
    print(
        f"[mask] stored composite at "
        f"input/{orig_subfolder}/{orig_filename} AND input/clipspace/{orig_filename} -> {url}"
    )
    return {
        "name": orig_filename,
        "subfolder": orig_subfolder,
        "type": "input",
        "url": url,
    }


@app.post("/api/upload/image")
async def upload_image(
    image: UploadFile = File(...),
    subfolder: str = Form(""),
    file_type: str = Form("input", alias="type"),
    overwrite: str = Form("false"),
):
    MAX_FILE_SIZE = 500 * 1024 * 1024  # 500MB

    content = await image.read()
    if len(content) > MAX_FILE_SIZE:
        raise HTTPException(413, "File too large")

    # Normalize filename (macOS screenshots use narrow no-break space \u202f)
    import unicodedata
    filename = unicodedata.normalize("NFKC", image.filename or "output").replace("\u202f", " ").replace("\u00a0", " ")
    subfolder = (subfolder or "").strip().strip("/")
    file_type = (file_type or "input").strip()

    content_type = image.content_type or "application/octet-stream"
    url = await fal_client.upload_async(content, content_type=content_type)

    # Also save to S3 if configured
    s3_key = await s3_upload(content, filename, content_type, "inputs")

    key = FileStorage.get_key(filename, subfolder, file_type)
    file_storage[key] = url
    return {"name": filename, "subfolder": subfolder, "type": file_type, "url": url, "s3_key": s3_key}


# --- S3 browsing endpoints ---

@app.get("/api/s3/status")
async def s3_status():
    """Check if S3 is configured and working."""
    if not s3_client:
        return {"enabled": False, "message": "S3 not configured. Set S3_BUCKET in .env"}
    try:
        s3_client.head_bucket(Bucket=S3_BUCKET)
        return {"enabled": True, "bucket": S3_BUCKET, "prefix": S3_PREFIX, "region": S3_REGION}
    except Exception as e:
        return {"enabled": False, "error": str(e)}


@app.get("/api/s3/list")
async def s3_list_files(folder: str = "outputs", max_keys: int = 100):
    """List files stored in S3."""
    if not s3_client:
        return {"files": [], "message": "S3 not configured"}
    try:
        prefix = f"{S3_PREFIX}/{folder}/"
        response = s3_client.list_objects_v2(Bucket=S3_BUCKET, Prefix=prefix, MaxKeys=max_keys)
        files = []
        for obj in response.get("Contents", []):
            filename = obj["Key"].removeprefix(prefix)
            files.append({
                "key": obj["Key"],
                "filename": filename,
                "size": obj["Size"],
                "last_modified": obj["LastModified"].isoformat(),
                "url": s3_get_url(obj["Key"]),
            })
        return {"files": files, "count": len(files)}
    except Exception as e:
        raise HTTPException(500, detail=f"S3 list failed: {str(e)}")


@app.get("/api/s3/download/{folder}/{filename:path}")
async def s3_download(folder: str, filename: str):
    """Get a presigned download URL for an S3 file."""
    if not s3_client:
        raise HTTPException(400, detail="S3 not configured")
    key = f"{S3_PREFIX}/{folder}/{filename}"
    try:
        url = s3_get_url(key)
        return {"url": url, "key": key}
    except Exception as e:
        raise HTTPException(500, detail=f"S3 download failed: {str(e)}")


# --- Frontend static file serving ---

@app.get("/materialdesignicons.min.css")
async def serve_mdi_css():
    if _frontend_css and os.path.isfile(_frontend_css):
        return Response(
            content=open(_frontend_css, "rb").read(),
            media_type="text/css",
        )
    raise HTTPException(404)


@app.get("/")
async def serve_index():
    index_path = os.path.join(FRONTEND_DIR, "index.html")
    if os.path.isfile(index_path):
        return Response(
            content=open(index_path, "rb").read(),
            media_type="text/html",
        )
    return {"message": "ComfyUI Proxy running. Frontend build not found."}


# Catch-all for SPA routing — must be LAST
@app.get("/{path:path}")
async def serve_frontend_catchall(path: str):
    # First check if it's a real file in the frontend build
    file_path = os.path.join(FRONTEND_DIR, path)
    if os.path.isfile(file_path):
        # Determine content type
        ext = os.path.splitext(path)[1].lower()
        content_types = {
            ".js": "application/javascript",
            ".css": "text/css",
            ".html": "text/html",
            ".json": "application/json",
            ".png": "image/png",
            ".svg": "image/svg+xml",
            ".ico": "image/x-icon",
            ".woff": "font/woff",
            ".woff2": "font/woff2",
            ".ttf": "font/ttf",
        }
        return Response(
            content=open(file_path, "rb").read(),
            media_type=content_types.get(ext, "application/octet-stream"),
        )

    # Only SPA fallback for routes WITHOUT file extensions (not .js, .css, .json, etc.)
    ext = os.path.splitext(path)[1].lower()
    if ext and ext != ".html":
        # This is a file request that wasn't found — return 404, NOT index.html
        raise HTTPException(404, detail=f"File not found: {path}")

    # SPA fallback for navigation routes
    index_path = os.path.join(FRONTEND_DIR, "index.html")
    if os.path.isfile(index_path):
        return Response(
            content=open(index_path, "rb").read(),
            media_type="text/html",
        )

    raise HTTPException(404)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8188)
