# /root/code/python/ieops-python/sdk/src/handler/server/comfyui.py
"""
ComfyUI Server - ComfyUI 代理服务接口层

设计模式与 OpenAIServer/DiffuserServer 一致：
- Server 只负责 HTTP 代理接口层
- 子进程管理由 handler 中的 ComfyUIProcess 实现

架构：
┌─────────────────────────────────────────────────────────────┐
│  ComfyUIServer (SDK - 代理接口层)                            │
│  - HTTP/WebSocket 代理                                      │
│  - 请求解析、响应格式化                                      │
│  - 接收 process_manager 作为参数                            │
└─────────────────────────────────────────────────────────────┘
                              ↓ 代理请求
┌─────────────────────────────────────────────────────────────┐
│  ComfyUIProcess (Worker handler.py)                         │
│  - 子进程启动和管理                                          │
│  - 健康检查                                                 │
└─────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────┐
│  ComfyUI 进程 (子进程)                                       │
└─────────────────────────────────────────────────────────────┘
"""

import asyncio
import json
import os
import base64
import uuid
import time
from typing import Any, AsyncGenerator, Dict, List, Optional, Callable, Protocol
from typing_extensions import TypedDict

import aiohttp
import uvicorn
from fastapi import FastAPI, Request, Response, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse, JSONResponse
from starlette.websockets import WebSocketState
from contextlib import asynccontextmanager

from .server import BaseServer
from ...utils.log import uvicorn_logger, log
from ...utils.util import APP_ID, arandstr

# ========== 环境变量配置 ==========
_MODEL_SERVER_SOCKET = os.getenv("MODEL_SERVER_SOCKET", f"{APP_ID}.sock")
_MODEL_SERVER_PORT = int(os.getenv("MODEL_SERVER_PORT", 8001))
_MODEL_SERVER_HOST = os.getenv("MODEL_SERVER_HOST", "127.0.0.1")
_GRACEFUL_SHUTDOWN_TIME = int(os.getenv("GRACEFUL_SHUTDOWN_TIME", 3))
_MODEL_TIMEOUT = int(os.getenv("MODEL_TIMEOUT", 600))

SERVER_CODE_OK = 0
SERVER_CODE_ERROR = 1
SERVER_CODE_STOP = 2


class ComfyUITask(TypedDict):
    prompt_id: str
    trace_id: str
    status: str
    progress: float
    images: List[str]


class ProcessManagerProtocol(Protocol):
    """进程管理器协议（用于类型提示）"""
    @property
    def is_running(self) -> bool: ...
    @property
    def is_ready(self) -> bool: ...
    async def start(self) -> None: ...
    async def stop(self) -> None: ...


class ComfyUIServer(BaseServer):
    """
    ComfyUI 代理服务器 - 只负责 HTTP/WebSocket 代理接口层
    
    与 OpenAIServer/DiffuserServer 设计模式一致：
    - 接收 process_manager 作为参数（由 handler 创建和管理）
    - 只负责 HTTP 代理、请求解析、响应格式化
    
    使用示例:
    ```python
    # handler.py
    class ComfyUIProcess:
        def __init__(self):
            self._manager = SubprocessManager(...)
        
        async def start(self):
            await self._manager.start()
        
        async def stop(self):
            await self._manager.stop()
        
        @property
        def is_running(self) -> bool:
            return self._manager.is_running
    
    process = ComfyUIProcess()
    server = ComfyUIServer(
        process_manager=process,
        comfyui_url="http://127.0.0.1:8188",
        on_workflow_received=preprocess_workflow,
    )
    Handler(server=server).serve()
    ```
    """
    
    def __init__(
        self,
        process_manager: Optional[ProcessManagerProtocol] = None,
        comfyui_url: Optional[str] = None,
        comfyui_host: Optional[str] = None,
        comfyui_port: Optional[int] = None,
        on_workflow_received: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
    ) -> None:
        """
        Args:
            process_manager: 进程管理器（由 handler 创建）
            comfyui_url: ComfyUI 完整 URL（优先使用）
            comfyui_host: ComfyUI 主机地址
            comfyui_port: ComfyUI 端口
            on_workflow_received: 工作流预处理回调
        """
        # ComfyUI URL
        if comfyui_url:
            self._comfyui_url = comfyui_url
        else:
            host = comfyui_host or os.getenv("COMFYUI_HOST", "127.0.0.1")
            port = comfyui_port or int(os.getenv("COMFYUI_PORT", 8188))
            self._comfyui_url = f"http://{host}:{port}"
        
        # 解析 host 和 port（用于 WebSocket）
        url_parts = self._comfyui_url.replace("http://", "").replace("https://", "").split(":")
        self._comfyui_host = url_parts[0]
        self._comfyui_port = int(url_parts[1]) if len(url_parts) > 1 else 8188
        
        self._process_manager = process_manager
        self._on_workflow_received = on_workflow_received
        
        self._tasks: Dict[str, ComfyUITask] = {}
        self._http_session: Optional[aiohttp.ClientSession] = None
        self._app = FastAPI(lifespan=self._lifespan)
        self._server: Optional[uvicorn.Server] = None
        
        self._setup_routes()

    @asynccontextmanager
    async def _lifespan(self, app: FastAPI):
        """FastAPI 生命周期管理"""
        uvicorn_logger.info("ComfyUI Server starting...")
        
        # 1. 启动子进程（如果提供了 process_manager）
        if self._process_manager:
            await self._process_manager.start()
        
        # 2. 初始化 HTTP 客户端
        self._http_session = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=_MODEL_TIMEOUT)
        )
        
        yield
        
        # 清理
        uvicorn_logger.info("ComfyUI Server shutting down...")
        
        # 关闭 HTTP 客户端
        if self._http_session:
            await self._http_session.close()
        
        # 停止子进程
        if self._process_manager:
            await self._process_manager.stop()
        
        # 清理 socket 文件
        if _MODEL_SERVER_SOCKET and os.path.exists(_MODEL_SERVER_SOCKET):
            os.remove(_MODEL_SERVER_SOCKET)

    def shutdown(self) -> None:
        """请求服务器优雅关闭"""
        if self._server:
            self._server.should_exit = True

    async def serve(self) -> None:
        """启动代理服务器"""
        if _MODEL_SERVER_SOCKET:
            config = uvicorn.Config(
                self._app,
                uds=_MODEL_SERVER_SOCKET,
                log_level="info",
                loop="asyncio",
                timeout_graceful_shutdown=_GRACEFUL_SHUTDOWN_TIME
            )
        else:
            config = uvicorn.Config(
                self._app,
                host=_MODEL_SERVER_HOST,
                port=_MODEL_SERVER_PORT,
                log_level="info",
                loop="asyncio",
                timeout_graceful_shutdown=_GRACEFUL_SHUTDOWN_TIME
            )
        
        self._server = uvicorn.Server(config)
        await self._server.serve()

    def _setup_routes(self) -> None:
        """设置 API 路由"""
        
        # ========== 健康检查 ==========
        @self._app.get("/health")
        async def health():
            """健康检查接口"""
            if self._process_manager:
                if not self._process_manager.is_ready:
                    return JSONResponse({"status": "initializing"}, status_code=503)
                if not self._process_manager.is_running:
                    return JSONResponse(
                        {"status": "error", "message": "ComfyUI process died"},
                        status_code=503
                    )
            return {"status": "ok"}

        # ========== ComfyUI 原生 API - 直接代理 ==========
        
        @self._app.get("/")
        async def root():
            """代理 ComfyUI 根路径"""
            return await self._proxy_request("GET", "/")

        @self._app.get("/system_stats")
        async def system_stats():
            """系统状态"""
            return await self._proxy_request("GET", "/system_stats")

        @self._app.get("/queue")
        async def get_queue():
            """获取队列信息"""
            return await self._proxy_request("GET", "/queue")

        @self._app.post("/queue")
        async def post_queue(request: Request):
            """队列操作"""
            body = await request.body()
            return await self._proxy_request("POST", "/queue", body=body)

        @self._app.post("/prompt")
        async def prompt(request: Request):
            """提交工作流"""
            body = await request.body()
            
            if request.headers.get("X-Stream-Response", "").lower() == "true":
                return StreamingResponse(
                    self._sse_generator(body),
                    media_type="text/event-stream",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                        "X-Accel-Buffering": "no",
                    }
                )
            
            return await self._proxy_request("POST", "/prompt", body=body)

        @self._app.post("/interrupt")
        async def interrupt():
            """中断生成"""
            return await self._proxy_request("POST", "/interrupt")

        @self._app.get("/history")
        async def get_history():
            """获取历史记录"""
            return await self._proxy_request("GET", "/history")

        @self._app.get("/history/{prompt_id}")
        async def get_history_by_id(prompt_id: str):
            """获取指定历史"""
            return await self._proxy_request("GET", f"/history/{prompt_id}")

        @self._app.post("/history")
        async def post_history(request: Request):
            """清除历史"""
            body = await request.body()
            return await self._proxy_request("POST", "/history", body=body)

        @self._app.get("/view")
        async def view(request: Request):
            """查看图像"""
            return await self._proxy_request("GET", "/view", params=dict(request.query_params))

        @self._app.post("/upload/image")
        async def upload_image(request: Request):
            """上传图像"""
            body = await request.body()
            headers = {k: v for k, v in request.headers.items() 
                      if k.lower() in ("content-type", "content-length")}
            return await self._proxy_request("POST", "/upload/image", body=body, headers=headers)

        @self._app.post("/upload/mask")
        async def upload_mask(request: Request):
            """上传遮罩"""
            body = await request.body()
            headers = {k: v for k, v in request.headers.items() 
                      if k.lower() in ("content-type", "content-length")}
            return await self._proxy_request("POST", "/upload/mask", body=body, headers=headers)

        @self._app.get("/embeddings")
        async def get_embeddings():
            """获取 embeddings"""
            return await self._proxy_request("GET", "/embeddings")

        @self._app.get("/extensions")
        async def get_extensions():
            """获取扩展列表"""
            return await self._proxy_request("GET", "/extensions")

        @self._app.get("/object_info")
        async def get_object_info():
            """获取节点信息"""
            return await self._proxy_request("GET", "/object_info")

        @self._app.get("/object_info/{node_class}")
        async def get_object_info_by_class(node_class: str):
            """获取指定节点信息"""
            return await self._proxy_request("GET", f"/object_info/{node_class}")

        # ========== 扩展接口 ==========
        
        @self._app.post("/prompt/stream")
        async def prompt_stream(request: Request):
            """流式工作流"""
            body = await request.body()
            return StreamingResponse(
                self._sse_generator(body),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )

        @self._app.post("/prompt/sync")
        async def prompt_sync(request: Request):
            """同步工作流"""
            body = await request.body()
            results: List[Any] = []
            last_code = SERVER_CODE_OK
            
            async for item in self._generate(body):
                code = item["code"]
                data = item["data"]
                if code == SERVER_CODE_ERROR:
                    last_code = code
                    results.append(data)
                    break
                if code == SERVER_CODE_STOP:
                    break
                if data is not None:
                    results.append(data)
            
            return {"code": last_code, "results": results}

        # ========== ieops 协议兼容 ==========
        @self._app.post("/sse/infer")
        async def infer_sse(request: Request):
            """SSE 推理接口"""
            body = await request.body()
            return StreamingResponse(
                self._sse_generator(body),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )

        @self._app.post("/infer")
        async def infer(request: Request):
            """同步推理接口"""
            body = await request.body()
            results: List[Any] = []
            last_code = SERVER_CODE_OK
            
            async for item in self._generate(body):
                code = item["code"]
                data = item["data"]
                if code == SERVER_CODE_ERROR:
                    last_code = code
                    results.append(data)
                    break
                if code == SERVER_CODE_STOP:
                    break
                if data is not None:
                    results.append(data)
            
            return {"code": last_code, "results": results}

        # ========== WebSocket 代理 ==========
        @self._app.websocket("/ws")
        async def websocket_proxy(websocket: WebSocket):
            """代理 WebSocket"""
            await self._proxy_websocket(websocket)
        
        # ========== 通配符路由 ==========
        @self._app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
        async def proxy_fallback(request: Request, path: str):
            """透传其他路由"""
            body = await request.body() if request.method in ["POST", "PUT"] else None
            return await self._proxy_request(
                request.method, 
                f"/{path}", 
                body=body,
                params=dict(request.query_params)
            )

    async def _proxy_request(
        self, 
        method: str, 
        path: str, 
        body: Optional[bytes] = None,
        params: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None
    ) -> Response:
        """通用代理请求"""
        url = f"{self._comfyui_url}{path}"
        
        try:
            async with self._http_session.request(
                method=method,
                url=url,
                data=body,
                params=params,
                headers=headers,
            ) as resp:
                content = await resp.read()
                
                response_headers = {}
                for k, v in resp.headers.items():
                    if k.lower() not in ("transfer-encoding", "content-encoding", "content-length"):
                        response_headers[k] = v
                
                return Response(
                    content=content,
                    status_code=resp.status,
                    headers=response_headers,
                    media_type=resp.content_type
                )
        except Exception as e:
            uvicorn_logger.error(f"Proxy error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    async def _proxy_websocket(self, websocket: WebSocket) -> None:
        """代理 WebSocket"""
        query_string = str(websocket.scope.get("query_string", b""), "utf-8")
        ws_url = f"ws://{self._comfyui_host}:{self._comfyui_port}/ws"
        if query_string:
            ws_url += f"?{query_string}"
        
        uvicorn_logger.info(f"WebSocket proxy connecting to: {ws_url}")
        should_stop = asyncio.Event()
        
        try:
            await websocket.accept()
            
            async with aiohttp.ClientSession() as ws_session:
                async with ws_session.ws_connect(ws_url, heartbeat=30.0, receive_timeout=None) as backend_ws:
                    uvicorn_logger.info("WebSocket proxy connected to ComfyUI")
                    
                    async def forward_to_backend():
                        try:
                            while not should_stop.is_set():
                                try:
                                    data = await asyncio.wait_for(websocket.receive(), timeout=5.0)
                                    if data["type"] == "websocket.receive":
                                        if "text" in data:
                                            await backend_ws.send_str(data["text"])
                                        elif "bytes" in data:
                                            await backend_ws.send_bytes(data["bytes"])
                                    elif data["type"] == "websocket.disconnect":
                                        should_stop.set()
                                        break
                                except asyncio.TimeoutError:
                                    continue
                        except WebSocketDisconnect:
                            should_stop.set()
                        except Exception as e:
                            uvicorn_logger.error(f"Forward to backend error: {e}")
                            should_stop.set()
                    
                    async def forward_to_client():
                        try:
                            while not should_stop.is_set():
                                try:
                                    msg = await asyncio.wait_for(backend_ws.receive(), timeout=5.0)
                                    if websocket.client_state != WebSocketState.CONNECTED:
                                        should_stop.set()
                                        break
                                    if msg.type == aiohttp.WSMsgType.TEXT:
                                        await websocket.send_text(msg.data)
                                    elif msg.type == aiohttp.WSMsgType.BINARY:
                                        await websocket.send_bytes(msg.data)
                                    elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                                        should_stop.set()
                                        break
                                except asyncio.TimeoutError:
                                    continue
                        except Exception as e:
                            uvicorn_logger.error(f"Forward to client error: {e}")
                            should_stop.set()
                    
                    tasks = [
                        asyncio.create_task(forward_to_backend()),
                        asyncio.create_task(forward_to_client()),
                    ]
                    
                    try:
                        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
                        should_stop.set()
                        for task in pending:
                            task.cancel()
                            try:
                                await task
                            except asyncio.CancelledError:
                                pass
                    except Exception:
                        should_stop.set()
                        for task in tasks:
                            task.cancel()
                
        except aiohttp.ClientError as e:
            uvicorn_logger.error(f"WebSocket proxy connection error: {e}")
            if websocket.client_state == WebSocketState.CONNECTED:
                await websocket.close(code=1011, reason="Backend connection failed")
        except Exception as e:
            uvicorn_logger.error(f"WebSocket proxy error: {e}")
            if websocket.client_state == WebSocketState.CONNECTED:
                await websocket.close(code=1011, reason=f"Proxy error: {e}")
        finally:
            uvicorn_logger.info("WebSocket proxy connection closed")

    async def _sse_generator(self, request: bytes) -> AsyncGenerator[str, None]:
        """SSE 生成器"""
        async for item in self._generate(request):
            # 直接序列化整个 item，data 已经是原始数据结构
            yield f"{json.dumps(item, ensure_ascii=False)}\n\n"

    async def _generate(self, request: bytes) -> AsyncGenerator[Dict[str, Any], None]:
        """核心生成逻辑"""
        start = int(time.time() * 1000)
        
        try:
            request_data = json.loads(request.decode("utf-8"))
            headers = request_data.get("headers", {})
            payload = request_data.get("payload", request_data)
        except json.JSONDecodeError:
            payload = json.loads(request.decode("utf-8"))
            headers = {}
        
        trace_id = str(headers.get("x-trace-id", f"traceid-{await arandstr(8)}"))
        logger = log.get_logger(trace_id=trace_id)
        logger.info("ComfyUI request received")
        
        workflow: Dict[str, Any] = payload if isinstance(payload, dict) else json.loads(payload)
        
        if self._on_workflow_received:
            workflow = self._on_workflow_received(workflow)
        
        client_id = str(uuid.uuid4())
        
        try:
            prompt_id = await self._queue_prompt(workflow, client_id)
            if not prompt_id:
                yield {"code": SERVER_CODE_ERROR, "data": {"error": "Failed to queue prompt"}}
                return
            
            logger.info(f"Prompt queued: {prompt_id}")
            
            self._tasks[prompt_id] = ComfyUITask(
                prompt_id=prompt_id,
                trace_id=trace_id,
                status="pending",
                progress=0.0,
                images=[]
            )
            
            ws_url = f"ws://{self._comfyui_host}:{self._comfyui_port}/ws?clientId={client_id}"
            
            async with self._http_session.ws_connect(ws_url) as ws:
                async for msg in ws:
                    if msg.type == aiohttp.WSMsgType.TEXT:
                        data = json.loads(msg.data)
                        msg_type = data.get("type")
                        msg_data = data.get("data", {})
                        
                        if msg_type == "progress":
                            progress = msg_data.get("value", 0) / max(msg_data.get("max", 1), 1)
                            self._tasks[prompt_id]["progress"] = progress
                            self._tasks[prompt_id]["status"] = "running"
                            yield {
                                "code": SERVER_CODE_OK,
                                "data": {
                                    "type": "progress",
                                    "progress": progress,
                                    "step": msg_data.get("value"),
                                    "total_steps": msg_data.get("max"),
                                    "node": msg_data.get("node")
                                }
                            }
                        
                        elif msg_type == "executing":
                            node = msg_data.get("node")
                            if node is None and msg_data.get("prompt_id") == prompt_id:
                                self._tasks[prompt_id]["status"] = "completed"
                                break
                            else:
                                yield {
                                    "code": SERVER_CODE_OK,
                                    "data": {"type": "executing", "node": node}
                                }
                        
                        elif msg_type == "executed":
                            if msg_data.get("prompt_id") == prompt_id:
                                output = msg_data.get("output", {})
                                if "images" in output:
                                    yield {
                                        "code": SERVER_CODE_OK,
                                        "data": {
                                            "type": "node_output",
                                            "node": msg_data.get("node"),
                                            "images_info": output["images"]
                                        }
                                    }
                        
                        elif msg_type == "execution_error":
                            error_msg = msg_data.get("exception_message", "Unknown error")
                            logger.error(f"Execution error: {error_msg}")
                            self._tasks[prompt_id]["status"] = "error"
                            yield {"code": SERVER_CODE_ERROR, "data": {"error": error_msg}}
                            return
                    
                    elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
                        yield {"code": SERVER_CODE_ERROR, "data": {"error": "WebSocket error"}}
                        return
            
            images = await self._get_images(prompt_id)
            self._tasks[prompt_id]["images"] = images
            
            yield {
                "code": SERVER_CODE_OK,
                "data": {
                    "type": "result",
                    "images": images,
                    "prompt_id": prompt_id
                }
            }
            yield {"code": SERVER_CODE_STOP, "data": None}
            
        except asyncio.TimeoutError:
            logger.error("Generation timeout")
            yield {"code": SERVER_CODE_ERROR, "data": {"error": "timeout"}}
        except Exception as e:
            logger.error(f"Error: {e}")
            yield {"code": SERVER_CODE_ERROR, "data": {"error": str(e)}}
        finally:
            cost = int(time.time() * 1000) - start
            logger.info(f"Generation completed, cost: {cost}ms")

    async def _queue_prompt(self, workflow: Dict[str, Any], client_id: str) -> Optional[str]:
        """提交工作流"""
        try:
            async with self._http_session.post(
                f"{self._comfyui_url}/prompt",
                json={"prompt": workflow, "client_id": client_id}
            ) as resp:
                if resp.status == 200:
                    result = await resp.json()
                    return result.get("prompt_id")
                else:
                    error = await resp.text()
                    uvicorn_logger.error(f"Failed to queue prompt: {error}")
                    return None
        except Exception as e:
            uvicorn_logger.error(f"Error queuing prompt: {e}")
            return None

    async def _get_images(self, prompt_id: str) -> List[str]:
        """获取图像"""
        images = []
        try:
            async with self._http_session.get(f"{self._comfyui_url}/history/{prompt_id}") as resp:
                if resp.status == 200:
                    history = await resp.json()
                    if prompt_id in history:
                        outputs = history[prompt_id].get("outputs", {})
                        for node_id, node_output in outputs.items():
                            if "images" in node_output:
                                for img_info in node_output["images"]:
                                    params = {
                                        "filename": img_info["filename"],
                                        "subfolder": img_info.get("subfolder", ""),
                                        "type": img_info.get("type", "output")
                                    }
                                    async with self._http_session.get(
                                        f"{self._comfyui_url}/view",
                                        params=params
                                    ) as img_resp:
                                        if img_resp.status == 200:
                                            img_data = await img_resp.read()
                                            images.append(base64.b64encode(img_data).decode("utf-8"))
        except Exception as e:
            uvicorn_logger.error(f"Error getting images: {e}")
        return images
