from __future__ import annotations

import asyncio
import json
import logging
import mimetypes
import os
from collections.abc import Callable
from pathlib import Path
from typing import Any
from urllib.parse import quote

import httpx

try:
    from websockets.asyncio.client import connect as ws_connect
except ImportError:  # websockets < 13
    from websockets import connect as ws_connect

from websockets.exceptions import ConnectionClosed

logger = logging.getLogger(__name__)

ProgressCallback = Callable[[int], None]


class ComfyError(RuntimeError):
    user_message = "Ошибка при обращении к ComfyUI"


class ComfyConnectionError(ComfyError):
    user_message = "Нет связи с ComfyUI"


class ComfyHTTPError(ComfyError):
    user_message = "ComfyUI вернул ошибку"


class ComfyValidationError(ComfyError):
    user_message = "Workflow отклонён ComfyUI"


class ComfyExecutionError(ComfyError):
    user_message = "Генерация завершилась с ошибкой"


class ComfyTimeout(ComfyError):
    user_message = "Генерация не завершилась за отведённое время"


def _format_api_error(payload: dict[str, Any]) -> str:
    parts: list[str] = []
    error = payload.get("error")
    if isinstance(error, dict):
        message = error.get("message") or error.get("type") or str(error)
        parts.append(str(message))
    elif error:
        parts.append(str(error))

    node_errors = payload.get("node_errors") or {}
    if isinstance(node_errors, dict):
        for node_id, info in node_errors.items():
            if not isinstance(info, dict):
                continue
            for item in info.get("errors") or []:
                if not isinstance(item, dict):
                    continue
                message = item.get("message") or item.get("type") or "ошибка"
                details = item.get("details")
                parts.append(f"узел {node_id}: {message}" + (f" ({details})" if details else ""))
    return "; ".join(parts) or "workflow отклонён"


def _execution_error_message(entry: dict[str, Any]) -> str:
    status = entry.get("status") or {}
    for message in status.get("messages") or []:
        if not isinstance(message, (list, tuple)) or len(message) != 2:
            continue
        event, data = message
        if event == "execution_error" and isinstance(data, dict):
            text = data.get("exception_message") or data.get("exception_type")
            if text:
                return str(text)
    return "ComfyUI сообщил об ошибке выполнения"


class ProgressSession:
    def __init__(self, websocket: Any) -> None:
        self._ws = websocket
        self._closed = False

    async def wait(
        self,
        prompt_id: str,
        *,
        timeout: float,
        on_progress: ProgressCallback,
    ) -> bool:
        """True — выполнение завершилось; False — WS потерян, нужен polling."""
        loop = asyncio.get_running_loop()
        deadline = loop.time() + timeout

        try:
            while True:
                remaining = deadline - loop.time()
                if remaining <= 0:
                    raise ComfyTimeout("WebSocket: превышено время ожидания")
                try:
                    raw = await asyncio.wait_for(self._ws.recv(), timeout=remaining)
                except TimeoutError as exc:
                    raise ComfyTimeout("WebSocket: превышено время ожидания") from exc
                except ConnectionClosed:
                    logger.warning("WebSocket ComfyUI закрыт, перехожу на polling")
                    return False

                if isinstance(raw, bytes):
                    raw = raw.decode("utf-8", "replace")
                try:
                    message = json.loads(raw)
                except json.JSONDecodeError:
                    continue
                if not isinstance(message, dict):
                    continue

                message_type = message.get("type")
                data = message.get("data") or {}
                if not isinstance(data, dict):
                    data = {}

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

                if message_type == "progress":
                    value = data.get("value", 0)
                    maximum = data.get("max", 0)
                    if isinstance(value, (int, float)) and isinstance(maximum, (int, float)) and maximum > 0:
                        on_progress(max(0, min(100, int(value * 100 / maximum))))
                elif message_type == "execution_error":
                    raise ComfyExecutionError(_execution_error_message({"status": {"messages": [["execution_error", data]]}}))
                elif message_type == "execution_interrupted":
                    raise ComfyExecutionError("Генерация прервана")
                elif message_type == "execution_success" or (
                    message_type == "executing" and data.get("node") is None
                ):
                    return True
        except ConnectionClosed:
            logger.warning("WebSocket ComfyUI закрыт, перехожу на polling")
            return False

    async def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        try:
            await self._ws.close()
        except Exception:
            logger.debug("Ошибка закрытия WebSocket", exc_info=True)


class ComfyClient:
    def __init__(self, base_url: str, ws_url: str) -> None:
        self._base_url = base_url.rstrip("/")
        self._ws_url = ws_url
        self._client = httpx.AsyncClient(
            base_url=self._base_url,
            timeout=httpx.Timeout(30.0, connect=10.0),
            trust_env=False,
        )

    async def close(self) -> None:
        await self._client.aclose()

    async def submit(self, workflow: dict[str, Any], client_id: str) -> str:
        try:
            response = await self._client.post(
                "/prompt",
                json={"prompt": workflow, "client_id": client_id},
            )
        except httpx.HTTPError as exc:
            raise ComfyConnectionError(f"POST /prompt: {exc}") from exc

        if response.status_code >= 400:
            raise ComfyHTTPError(
                f"POST /prompt: HTTP {response.status_code}: {response.text[:500]}"
            )
        try:
            payload = response.json()
        except ValueError as exc:
            raise ComfyHTTPError("POST /prompt: ответ не является JSON") from exc

        if payload.get("error") or payload.get("node_errors"):
            raise ComfyValidationError(_format_api_error(payload))

        prompt_id = payload.get("prompt_id")
        if not prompt_id:
            raise ComfyHTTPError("POST /prompt: в ответе нет prompt_id")
        return str(prompt_id)

    async def upload_image(self, path: Path, remote_name: str) -> str:
        try:
            content = path.read_bytes()
        except OSError as exc:
            raise ComfyError(f"Не удалось прочитать входной файл {path}: {exc}") from exc

        mime = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
        try:
            response = await self._client.post(
                "/upload/image",
                files={"image": (remote_name, content, mime)},
                data={"type": "input", "overwrite": "true"},
                timeout=60.0,
            )
        except httpx.HTTPError as exc:
            raise ComfyConnectionError(f"POST /upload/image: {exc}") from exc

        if response.status_code >= 400:
            raise ComfyHTTPError(
                f"POST /upload/image: HTTP {response.status_code}: {response.text[:500]}"
            )
        try:
            payload = response.json()
        except ValueError as exc:
            raise ComfyHTTPError("POST /upload/image: ответ не является JSON") from exc

        name = payload.get("name")
        if not name:
            raise ComfyHTTPError("POST /upload/image: в ответе нет имени файла")
        subfolder = payload.get("subfolder") or ""
        return f"{subfolder}/{name}" if subfolder else str(name)

    async def get_history(self, prompt_id: str) -> dict[str, Any] | None:
        try:
            response = await self._client.get(f"/history/{quote(prompt_id)}")
        except httpx.HTTPError as exc:
            raise ComfyConnectionError(f"GET /history: {exc}") from exc

        if response.status_code >= 400:
            raise ComfyHTTPError(
                f"GET /history: HTTP {response.status_code}: {response.text[:500]}"
            )
        try:
            payload = response.json()
        except ValueError as exc:
            raise ComfyHTTPError("GET /history: ответ не является JSON") from exc
        entry = payload.get(prompt_id)
        return entry if isinstance(entry, dict) else None

    async def wait_for_history(self, prompt_id: str, *, timeout: float) -> dict[str, Any]:
        loop = asyncio.get_running_loop()
        deadline = loop.time() + timeout
        while loop.time() < deadline:
            entry = await self.get_history(prompt_id)
            if entry is not None:
                status = entry.get("status") or {}
                if status.get("status_str") == "error":
                    raise ComfyExecutionError(_execution_error_message(entry))
                if status.get("completed") or status.get("status_str") == "success":
                    return entry
            await asyncio.sleep(min(2.0, max(0.1, deadline - loop.time())))
        raise ComfyTimeout("Не дождались результата в истории ComfyUI")

    def find_output_image(self, entry: dict[str, Any]) -> dict[str, Any] | None:
        outputs = entry.get("outputs") or {}
        if not isinstance(outputs, dict):
            return None
        for node_output in outputs.values():
            if not isinstance(node_output, dict):
                continue
            for image in node_output.get("images") or []:
                if isinstance(image, dict) and image.get("filename"):
                    return image
        return None

    async def download_output(self, entry: dict[str, Any], destination: Path) -> Path:
        image = self.find_output_image(entry)
        if image is None:
            raise ComfyError("ComfyUI не вернул изображение")
        return await self.download_image(
            filename=str(image["filename"]),
            subfolder=str(image.get("subfolder") or ""),
            image_type=str(image.get("type") or "output"),
            destination=destination,
        )

    async def download_image(
        self,
        *,
        filename: str,
        subfolder: str,
        image_type: str,
        destination: Path,
    ) -> Path:
        destination.parent.mkdir(parents=True, exist_ok=True)
        temp_path = destination.with_name(destination.name + ".part")
        params = {"filename": filename, "subfolder": subfolder, "type": image_type}
        try:
            async with self._client.stream(
                "GET", "/view", params=params, timeout=httpx.Timeout(120.0, connect=10.0)
            ) as response:
                if response.status_code >= 400:
                    body = (await response.aread())[:500]
                    raise ComfyHTTPError(
                        f"GET /view: HTTP {response.status_code}: {body!r}"
                    )
                with temp_path.open("wb") as handle:
                    async for chunk in response.aiter_bytes():
                        handle.write(chunk)
            os.replace(temp_path, destination)
        except httpx.HTTPError as exc:
            temp_path.unlink(missing_ok=True)
            raise ComfyConnectionError(f"GET /view: {exc}") from exc
        except Exception:
            temp_path.unlink(missing_ok=True)
            raise
        return destination

    async def open_progress(self, client_id: str) -> ProgressSession | None:
        url = f"{self._ws_url}?clientId={quote(client_id)}"
        try:
            websocket = await ws_connect(
                url,
                open_timeout=5,
                close_timeout=2,
                max_size=8 * 1024 * 1024,
                ping_interval=20,
            )
        except Exception as exc:  # noqa: BLE001 - WS не должен ронять генерацию
            logger.warning("Не удалось подключиться к WebSocket ComfyUI: %s", exc)
            return None
        return ProgressSession(websocket)