"""
title: Comfier
author: Felix Morgner
author_url: https://felixmorgner.com
description: A ComfyUI image generation tool extending the built-in tool by allowing for more parameters.
required_open_webui_version: 0.12.0
version: 0.1.0
license: 3-clause-BSD
"""

import asyncio
import copy
import inspect
import json
import mimetypes
import random
import time
import uuid
from typing import Annotated, Any, Awaitable, Callable, Optional

import aiohttp
from fastapi import Request
from open_webui.models.config import Config
from open_webui.models.users import Users
from open_webui.routers.images import get_image_config, upload_image
from pydantic import BaseModel, Field

EventEmitter = Callable[[dict[str, Any]], Awaitable[None]]
EventCall = Callable[[dict[str, Any]], Awaitable[Any]]


class Tools:
    """
    Comfier: ComfyUI image generation for Open WebUI.

    Uses the workflow, connection settings and node mappings (prompt, seed, width,
    height, steps) configured in Open WebUI's image generation settings, and adds
    negative prompts, per-call size and reproducible seeds on top.
    """

    class Valves(BaseModel):
        negative_prompt: str = Field(
            default="watermarks",
            description="Default negative prompt, used when the model does not provide one.",
        )
        negative_prompt_key: str = Field(
            default="negative_prompt",
            description="Name of the workflow node input that receives the negative prompt (e.g. 'negative_prompt' or 'text').",
        )
        negative_prompt_node_ids: str = Field(
            default="",
            description="Comma-separated IDs of the workflow nodes that receive the negative prompt (e.g. '459:452'). Leave empty to disable negative prompting.",
        )

    def __init__(self):
        self.valves = self.Valves()

    async def comfier_generate_image(
        self,
        positive_prompt: str,
        negative_prompt: Optional[str] = None,
        size: Optional[Annotated[str, Field(pattern=r"^\d+[xX]\d+$")]] = None,
        seed: Optional[int | str] = None,
        __request__: Optional[Request] = None,
        __user__: Optional[dict[str, Any]] = None,
        __event_emitter__: Optional[EventEmitter] = None,
        __event_call__: Optional[EventCall] = None,
    ) -> dict[str, Any]:
        """
        Generate an image from a text prompt with ComfyUI.
        Optionally accepts a negative prompt (what to avoid), an image size and a seed for reproducible results.

        The generated image is automatically displayed to the user in the chat, so do not embed or link it again; just tell the user it has been created.
        Generation may take a while. The result contains the seed that was used (random if none was given).
        To reproduce an image, pass that seed again with the same prompt and size.

        :param positive_prompt: Descriptive text prompt for the image to generate.
        :param negative_prompt: Optional text describing what to avoid in the image. If omitted, the administrator's default negative prompt is used; pass an empty string for none.
        :param size: Optional image size in pixels as 'WIDTHxHEIGHT' (e.g. '1024x768'). If omitted, the default size configured in Open WebUI is used.
        :param seed: Optional non-negative integer seed (or its decimal string) for reproducible generation. If omitted, a random seed is used.
        :return: A dict with 'status', the 'seed' used, a 'message' for the model and the list of stored image URLs under 'images'.
        """

        if __request__ is None:
            raise RuntimeError(
                "FastAPI Request context (__request__) was not injected."
            )

        if not __user__ or "id" not in __user__:
            raise RuntimeError("User context (__user__) was not injected.")

        normalized_seed = self._normalize_seed(seed)
        workflow = await self._load_workflow()
        global_mappings = await self._load_global_node_mappings()
        effective_mappings = self._resolve_node_mappings(global_mappings)
        effective_negative_prompt = (
            negative_prompt
            if negative_prompt is not None
            else self.valves.negative_prompt
        )
        width, height = await self._resolve_dimensions(size)
        steps = await self._resolve_steps()

        parameter_values: dict[str, Any] = {
            "positive_prompt": positive_prompt,
            "negative_prompt": effective_negative_prompt,
            "seed": normalized_seed,
            "width": width,
            "height": height,
            "steps": steps,
        }

        self._apply_workflow_parameters(
            workflow=workflow,
            mappings=effective_mappings,
            values=parameter_values,
        )

        user = Users.get_user_by_id(__user__["id"])
        if inspect.isawaitable(user):
            user = await user

        while True:
            try:
                image_urls = await self._execute(
                    workflow,
                    request=__request__,
                    user=user,
                    metadata={
                        "prompt": positive_prompt,
                        "negative_prompt": effective_negative_prompt,
                        "seed": str(normalized_seed),
                        "width": width,
                        "height": height,
                        "steps": steps,
                    },
                    event_emitter=__event_emitter__,
                )
                break
            except (ConnectionError, aiohttp.ClientError) as err:
                if __event_call__ is None:
                    raise
                await self._emit_status(
                    __event_emitter__,
                    f"ComfyUI connection lost: {err}. Waiting for retry confirmation...",
                    done=False,
                )
                should_retry = await __event_call__(
                    {
                        "type": "confirmation",
                        "data": {
                            "title": "ComfyUI Backend Disconnected",
                            "message": (
                                f"Lost connection to ComfyUI during generation ({err}).\n\n"
                                f"Would you like to retry with the same parameters (seed {normalized_seed}, {width}x{height})?"
                            ),
                        },
                    }
                )
                if not should_retry:
                    await self._emit_status(
                        __event_emitter__,
                        f"Generation failed: {err}",
                        done=True,
                    )
                    raise

        if __event_emitter__ is not None and image_urls:
            await __event_emitter__(
                {
                    "type": "files",
                    "data": {
                        "files": [{"type": "image", "url": url} for url in image_urls]
                    },
                }
            )

        return {
            "status": "success",
            "seed": str(normalized_seed),
            "message": "The image has been successfully generated and is already visible to the user in the chat. You do not need to display or embed the image again - just acknowledge that it has been created.",
            "images": image_urls,
        }

    async def comfier_generate_seed(
        self, bits: Annotated[int, Field(ge=8, le=63)] = 63
    ) -> str:
        """
        Generate a new random seed for image generation.

        Calling this is optional: comfier_generate_image picks a random seed itself and reports it in its result.
        Use this tool only when a seed is needed before generating, for example to create several images with distinct, known seeds.
        Returns the seed as a decimal string, which can be passed as the seed of comfier_generate_image.

        :param bits: Maximum bit-width of the seed, between 8 and 63. Larger values allow larger seeds. Defaults to 63.
        :return: The generated seed as a decimal string.
        """

        clamped = max(8, min(int(bits), 63))
        return f"{random.getrandbits(clamped)}"

    def _apply_workflow_parameters(
        self,
        workflow: dict[str, Any],
        mappings: dict[str, dict[str, Any]],
        values: dict[str, Any],
    ) -> dict[str, list[str]]:
        """
        Inject the resolved parameter values into the workflow dictionary in-place.

        Parameters without a mapping or with a value of None are skipped.
        Returns the injected targets ("node_id.input_key") per parameter name.
        Raises KeyError if a mapped node does not exist in the workflow.
        """

        applied: dict[str, list[str]] = {}

        for param_name, value in values.items():
            mapping = mappings.get(param_name)

            if not mapping or value is None:
                continue

            input_key = mapping.get("key")
            node_ids = mapping.get("node_ids", [])

            if not input_key or not node_ids:
                continue

            updated_targets: list[str] = []

            for node_id in node_ids:
                node = workflow.get(node_id)

                if not isinstance(node, dict) or not isinstance(
                    node.get("inputs"), dict
                ):
                    raise KeyError(
                        f"Mapped node {node_id!r} for {param_name!r} does not exist in COMFYUI_WORKFLOW."
                    )

                node["inputs"][input_key] = value
                updated_targets.append(f"{node_id}.{input_key}")

            if updated_targets:
                applied[param_name] = updated_targets

        return applied

    async def _cancel(
        self, http_base: str, headers: dict[str, str], prompt_id: Optional[str]
    ) -> None:
        """Try to cancel a ComfyUI request."""

        if not prompt_id:
            return

        try:
            timeout = aiohttp.ClientTimeout(total=3.0)
            async with aiohttp.ClientSession(
                headers=headers, timeout=timeout
            ) as cancel_session:
                await cancel_session.post(
                    f"{http_base}/queue", json={"delete": [prompt_id]}
                )
                await cancel_session.post(
                    f"{http_base}/interrupt", json={"prompt_id": prompt_id}
                )
        except Exception:
            pass

    async def _collect_outputs(
        self,
        ws: aiohttp.ClientWebSocketResponse,
        session: aiohttp.ClientSession,
        http_base: str,
        workflow: dict[str, Any],
        prompt_id: str,
        event_emitter: Optional[EventEmitter],
    ) -> list[dict[str, Any]]:
        """
        Follow the ComfyUI WebSocket until the prompt finishes, emitting progress.

        Returns the image descriptors reported by the workflow's output nodes.
        Raises ConnectionError if the backend disappears, RuntimeError if the
        job is interrupted or fails.
        """

        output_images: list[dict[str, Any]] = []
        last_emit_time = 0.0
        idle_ticks = 0

        while True:
            try:
                msg = await asyncio.wait_for(ws.receive(), timeout=1.0)
            except asyncio.TimeoutError:
                idle_ticks += 1
                if idle_ticks >= 5:
                    idle_ticks = 0
                    await self._ensure_prompt_alive(session, http_base, prompt_id)
                continue

            idle_ticks = 0

            if msg.type != aiohttp.WSMsgType.TEXT:
                self._raise_if_ws_closed(ws, msg.type)
                continue

            payload = json.loads(msg.data) or {}
            msg_type = payload.get("type")
            data = payload.get("data") or {}

            if data.get("prompt_id") and data.get("prompt_id") != prompt_id:
                continue

            if msg_type == "executing":
                if data.get("node") is None:
                    return output_images
                await self._emit_node_status(event_emitter, workflow, str(data["node"]))

            elif msg_type == "progress":
                last_emit_time = await self._emit_progress(
                    event_emitter, data, last_emit_time
                )

            elif msg_type == "executed":
                images = (data.get("output") or {}).get("images") or []
                if isinstance(images, list):
                    output_images.extend(img for img in images if isinstance(img, dict))

            elif msg_type == "execution_interrupted":
                raise RuntimeError("Generation was interrupted on the backend.")

            elif msg_type == "execution_error":
                error_message = data.get("exception_message", "Unknown ComfyUI error")
                raise RuntimeError(f"ComfyUI execution error: {error_message}")

    async def _emit_node_status(
        self,
        event_emitter: Optional[EventEmitter],
        workflow: dict[str, Any],
        node_id: str,
    ) -> None:
        """Announce the workflow node ComfyUI has started executing."""

        node_info = workflow.get(node_id) or {}
        node_title = (
            (node_info.get("_meta") or {}).get("title")
            or node_info.get("class_type")
            or f"Node {node_id}"
        )
        await self._emit_status(event_emitter, f"Running: {node_title}...", done=False)

    async def _emit_progress(
        self,
        event_emitter: Optional[EventEmitter],
        data: dict[str, Any],
        last_emit_time: float,
    ) -> float:
        """Emit a throttled sampling progress update. Returns the new last-emit time."""

        val = data.get("value", 0)
        maximum = data.get("max", 1)
        now = time.monotonic()

        if (now - last_emit_time) < 0.25 and val != maximum:
            return last_emit_time

        percentage = int((val / maximum) * 100) if maximum else 0
        await self._emit_status(
            event_emitter,
            f"Sampling: step {val}/{maximum} ({percentage}%)",
            done=False,
        )
        return now

    async def _emit_status(
        self,
        event_emitter: Optional[EventEmitter],
        description: str,
        done: bool = False,
    ) -> None:
        """Send a live status update banner to the Open WebUI frontend."""

        if event_emitter is None:
            return

        await event_emitter(
            {
                "type": "status",
                "data": {
                    "description": description,
                    "done": done,
                    "hidden": False,
                },
            }
        )

    async def _ensure_prompt_alive(
        self, session: aiohttp.ClientSession, http_base: str, prompt_id: str
    ) -> None:
        """Raise ConnectionError if ComfyUI no longer tracks the prompt."""

        if not await self._verify_prompt_alive(session, http_base, prompt_id):
            raise ConnectionError(
                f"ComfyUI backend lost prompt {prompt_id} (backend may have restarted or crashed)."
            )

    async def _execute(
        self,
        workflow: dict[str, Any],
        request: Request,
        user: Any,
        metadata: dict[str, Any],
        event_emitter: Optional[EventEmitter],
    ) -> list[str]:
        """Submit the workflow to ComfyUI, stream progress, and store the generated images."""

        http_base, api_key = await self._load_connection_information()
        ws_base = http_base.replace("https://", "wss://").replace("http://", "ws://")
        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
        client_id = str(uuid.uuid4())

        stored_urls: list[str] = []
        prompt_id: Optional[str] = None

        try:
            async with aiohttp.ClientSession(headers=headers) as session:
                async with session.ws_connect(
                    f"{ws_base}/ws?clientId={client_id}",
                    heartbeat=5.0,
                ) as ws:
                    await self._emit_status(
                        event_emitter, "Submitting workflow to ComfyUI...", done=False
                    )
                    prompt_id = await self._submit_workflow(
                        session, http_base, workflow, client_id
                    )
                    await self._emit_status(
                        event_emitter, "Waiting for ComfyUI queue...", done=False
                    )
                    output_images = await self._collect_outputs(
                        ws, session, http_base, workflow, prompt_id, event_emitter
                    )

                if not output_images:
                    raise RuntimeError(
                        "Generation finished without producing any output. The Job may have been cancelled or removed from the queue."
                    )

                await self._emit_status(
                    event_emitter, "Saving generated images...", done=False
                )

                for img in output_images:
                    url = await self._store_image(
                        session, http_base, img, request, user, metadata
                    )
                    if url:
                        stored_urls.append(url)

        except asyncio.CancelledError:
            await asyncio.shield(self._cancel(http_base, headers, prompt_id))
            await asyncio.shield(
                self._emit_status(
                    event_emitter, "Generation aborted by user", done=True
                )
            )
            raise

        except Exception as err:
            await self._emit_status(
                event_emitter, f"Generation failed: {err}", done=True
            )
            raise

        await self._emit_status(event_emitter, "Generation complete", done=True)
        return stored_urls

    async def _load_connection_information(self) -> tuple[str, str]:
        """Load COMFYUI_BASE_URL and COMFYUI_API_KEY from Open WebUI config."""

        image_cfg = await get_image_config()
        base_url = getattr(image_cfg, "COMFYUI_BASE_URL", None)
        api_key = getattr(image_cfg, "COMFYUI_API_KEY", None)

        if not base_url:
            base_url = await Config.get("image_generation.comfyui.base_url", "")

        if api_key is None:
            api_key = await Config.get("image_generation.comfyui.api_key", "")

        base_url = str(base_url or "").strip().rstrip("/")

        if not base_url:
            raise RuntimeError("COMFYUI_BASE_URL is not configured in Open WebUI.")

        return base_url, str(api_key or "").strip()

    async def _load_global_node_mappings(self) -> list[dict[str, Any]]:
        """Load the globally configured node mappings (COMFYUI_WORKFLOW_NODES) list."""

        image_cfg = await get_image_config()
        nodes = getattr(image_cfg, "COMFYUI_WORKFLOW_NODES", None)

        if nodes is None:
            nodes = await Config.get("image_generation.comfyui.nodes", [])

        return nodes if isinstance(nodes, list) else []

    async def _load_workflow(self) -> dict[str, Any]:
        """Load and parse the globally configured ComfyUI workflow JSON."""

        image_cfg = await get_image_config()
        raw_workflow = getattr(image_cfg, "COMFYUI_WORKFLOW", None)

        if not raw_workflow:
            raw_workflow = await Config.get("image_generation.comfyui.workflow", "")

        if isinstance(raw_workflow, dict):
            return copy.deepcopy(raw_workflow)

        if not isinstance(raw_workflow, str) or not raw_workflow.strip():
            raise RuntimeError(
                "No global COMFYUI_WORKFLOW is configured in Open WebUI."
            )

        try:
            workflow = json.loads(raw_workflow)
        except json.JSONDecodeError as err:
            raise RuntimeError(
                "Configured COMFYUI_WORKFLOW is not valid JSON."
            ) from err

        if not isinstance(workflow, dict):
            raise RuntimeError(
                "Configured COMFYUI_WORKFLOW root must be a JSON object."
            )

        return workflow

    def _normalize_seed(self, seed: Optional[int | str]) -> int:
        """Normalize an optional int or string seed into a valid 63-bit positive integer."""

        if seed is None:
            return random.getrandbits(63)

        if isinstance(seed, str):
            stripped = seed.strip()
            if not stripped:
                return random.getrandbits(63)
            try:
                parsed = int(stripped, 0)
            except ValueError as err:
                raise ValueError(f"Invalid seed string: {seed!r}") from err
        else:
            parsed = int(seed)

        if parsed < 0:
            return random.getrandbits(63)

        return parsed & ((1 << 63) - 1)

    def _raise_if_ws_closed(
        self, ws: aiohttp.ClientWebSocketResponse, msg_type: aiohttp.WSMsgType
    ) -> None:
        """Raise ConnectionError if a non-text WebSocket message signals a closed connection."""

        if msg_type not in (
            aiohttp.WSMsgType.CLOSE,
            aiohttp.WSMsgType.CLOSING,
            aiohttp.WSMsgType.CLOSED,
            aiohttp.WSMsgType.ERROR,
        ):
            return

        ws_err = ws.exception()
        detail = f": {ws_err}" if ws_err else f" (code {ws.close_code})"
        raise ConnectionError(
            f"ComfyUI WebSocket disconnected unexpectedly before completion{detail}."
        )

    async def _resolve_dimensions(self, size: Optional[str]) -> tuple[int, int]:
        """Parse 'WIDTHxHEIGHT' from the size parameter or fall back to global IMAGE_SIZE."""

        raw_size = (size or "").strip()

        if not raw_size:
            image_cfg = await get_image_config()
            raw_size = getattr(image_cfg, "IMAGE_SIZE", None)

            if not raw_size:
                raw_size = await Config.get("image_generation.size", "1024x1024")

        try:
            width_string, height_string = str(raw_size).lower().split("x", 1)
            width, height = int(width_string.strip()), int(height_string.strip())
        except (ValueError, AttributeError) as err:
            raise ValueError(
                f"Invalid image size {raw_size!r}. Expected 'WIDTHxHEIGHT' (e.g. '1024x768')."
            ) from err

        if width <= 0 or height <= 0:
            raise ValueError(f"Image dimensions must be positive: {width}x{height}")

        return width, height

    def _resolve_node_mappings(
        self, global_nodes: list[dict[str, Any]]
    ) -> dict[str, dict[str, Any]]:
        """Combine global prompt/seed mappings with the tool's negative_prompt Valve."""

        resolved: dict[str, dict[str, Any]] = {
            "positive_prompt": {"key": "text", "node_ids": []},
            "negative_prompt": {
                "key": self.valves.negative_prompt_key.strip() or "negative_prompt",
                "node_ids": [
                    nid.strip()
                    for nid in self.valves.negative_prompt_node_ids.split(",")
                    if nid.strip()
                ],
            },
            "seed": {"key": "seed", "node_ids": []},
            "width": {"key": "width", "node_ids": []},
            "height": {"key": "height", "node_ids": []},
            "steps": {"key": "steps", "node_ids": []},
        }

        type_map = {
            "prompt": "positive_prompt",
            "positive_prompt": "positive_prompt",
            "seed": "seed",
            "width": "width",
            "height": "height",
            "steps": "steps",
        }

        for entry in global_nodes:
            if not isinstance(entry, dict):
                continue

            target = type_map.get(entry.get("type", ""))

            if not target:
                continue

            if entry.get("key"):
                resolved[target]["key"] = str(entry["key"]).strip()

            if isinstance(entry.get("node_ids"), list):
                resolved[target]["node_ids"] = [
                    str(nid).strip() for nid in entry["node_ids"] if str(nid).strip()
                ]

        return resolved

    async def _resolve_steps(self) -> int:
        """Load the global IMAGE_STEPS setting from Open WebUI config."""

        image_cfg = await get_image_config()
        raw_steps = getattr(image_cfg, "IMAGE_STEPS", None)

        if raw_steps is None:
            raw_steps = await Config.get("image_generation.steps", 25)

        return int(raw_steps)

    async def _store_image(
        self,
        session: aiohttp.ClientSession,
        http_base: str,
        img: dict[str, Any],
        request: Request,
        user: Any,
        metadata: dict[str, Any],
    ) -> Optional[str]:
        """
        Download one generated image from ComfyUI and store it in Open WebUI.

        Returns the stored image's URL, or None if the descriptor has no filename.
        """

        filename = img.get("filename")
        if not filename:
            return None

        params = {
            "filename": filename,
            "subfolder": img.get("subfolder", ""),
            "type": img.get("type", "output"),
        }

        async with session.get(f"{http_base}/view", params=params) as img_response:
            img_response.raise_for_status()
            image_bytes = await img_response.read()
            content_type = (
                img_response.headers.get("Content-Type")
                or mimetypes.guess_type(filename)[0]
                or "image/png"
            )

        uploaded: Any = upload_image(
            request=request,
            image_data=image_bytes,
            content_type=content_type,
            metadata=metadata,
            user=user,
        )

        if inspect.isawaitable(uploaded):
            uploaded = await uploaded

        return self._upload_result_to_url(uploaded)

    async def _submit_workflow(
        self,
        session: aiohttp.ClientSession,
        http_base: str,
        workflow: dict[str, Any],
        client_id: str,
    ) -> str:
        """Queue the workflow on ComfyUI and return the new prompt ID."""

        async with session.post(
            f"{http_base}/prompt",
            json={"prompt": workflow, "client_id": client_id},
        ) as response:
            if response.status >= 400:
                error_body = await response.text()
                raise RuntimeError(
                    f"ComfyUI rejected workflow (HTTP {response.status}): {error_body}"
                )
            queued = await response.json()
            return queued["prompt_id"]

    def _upload_result_to_url(self, uploaded: Any) -> str:
        """Turn the return value of Open WebUI's upload_image into a file content URL."""

        if isinstance(uploaded, tuple):
            uploaded = uploaded[-1]

        if isinstance(uploaded, str):
            return uploaded

        if isinstance(uploaded, dict) and "id" in uploaded:
            return f"/api/v1/files/{uploaded['id']}/content"

        if hasattr(uploaded, "id"):
            return f"/api/v1/files/{uploaded.id}/content"  # pyright: ignore[reportAttributeAccessIssue]

        raise RuntimeError(
            f"Unexpected return type from upload_image: {type(uploaded)!r}"
        )

    async def _verify_prompt_alive(
        self,
        session: aiohttp.ClientSession,
        http_base: str,
        prompt_id: str,
    ) -> bool:
        """
        Verify with ComfyUI that the backend is reachable and still tracking prompt_id
        in either /queue (running/pending) or /history.
        """

        timeout = aiohttp.ClientTimeout(total=3.0)
        try:
            async with session.get(f"{http_base}/queue", timeout=timeout) as q_resp:
                q_resp.raise_for_status()
                q_data = await q_resp.json()

            running = q_data.get("queue_running") or []
            pending = q_data.get("queue_pending") or []
            in_queue = any(
                isinstance(item, list) and len(item) > 1 and item[1] == prompt_id
                for item in (running + pending)
            )
            if in_queue:
                return True

            async with session.get(
                f"{http_base}/history/{prompt_id}", timeout=timeout
            ) as h_resp:
                h_resp.raise_for_status()
                h_data = await h_resp.json() or {}
                if prompt_id not in h_data:
                    return False

                status_info = (h_data[prompt_id] or {}).get("status") or {}
                if status_info.get("status_str") == "error":
                    raise RuntimeError(
                        "The job was terminated with an error or interruption on the backend."
                    )
                return True

        except aiohttp.ClientError as err:
            raise ConnectionError(
                f"Lost connection to ComfyUI backend ({err})."
            ) from err