# /root/code/python/ieops-python/sdk/src/handler/server/comfyui.py

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

import aiohttp
import uvicorn
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse, JSONResponse
from contextlib import asynccontextmanager

from .server import BaseServer
from ...utils.log import *
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))

# ComfyUI 专用配置
_COMFYUI_PORT = int(os.getenv("COMFYUI_PORT", 8188))
_COMFYUI_HOST = os.getenv("COMFYUI_HOST", "127.0.0.1")
_COMFYUI_START_CMD = os.getenv("COMFYUI_START_CMD", "python main.py")
_COMFYUI_WORK_DIR = os.getenv("COMFYUI_WORK_DIR", "/workspace/ComfyUI")
_COMFYUI_STARTUP_TIMEOUT = int(os.getenv("COMFYUI_STARTUP_TIMEOUT", 120))

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 ComfyUIServer(BaseServer):
    """
    ComfyUI 代理服务器 - 管理独立 ComfyUI 进程的生命周期
    
    架构：
    ┌─────────────────┐     ┌──────────────────┐     ┌────────────────┐
    │  外部客户端      │────▶│  ComfyUIServer   │────▶│  ComfyUI 进程  │
    │  (HTTP/SSE)     │     │  (FastAPI 代理)   │     │  (子进程)      │
    └─────────────────┘     └──────────────────┘     └────────────────┘
    """
    
    def __init__(
        self,
        start_cmd: Optional[str] = None,
        work_dir: 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,
        extra_env: Optional[Dict[str, str]] = None,
    ) -> None:
        """
        Args:
            start_cmd: 启动 ComfyUI 的命令，如 "python main.py --port 8188"
            work_dir: ComfyUI 工作目录
            comfyui_host: ComfyUI 监听地址
            comfyui_port: ComfyUI 监听端口
            on_workflow_received: 工作流预处理回调
            extra_env: 额外的环境变量
        """
        self._start_cmd = start_cmd or _COMFYUI_START_CMD
        self._work_dir = work_dir or _COMFYUI_WORK_DIR
        self._comfyui_host = comfyui_host or _COMFYUI_HOST
        self._comfyui_port = comfyui_port or _COMFYUI_PORT
        self._comfyui_url = f"http://{self._comfyui_host}:{self._comfyui_port}"
        self._on_workflow_received = on_workflow_received
        self._extra_env = extra_env or {}
        
        self._comfyui_process: Optional[subprocess.Popen] = None
        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._comfyui_ready = False
        
        self._setup_routes()

    @asynccontextmanager
    async def _lifespan(self, app: FastAPI):
        """FastAPI 生命周期管理"""
        uvicorn_logger.info("ComfyUI Server starting...")
        
        # 1. 启动 ComfyUI 子进程
        await self._start_comfyui()
        
        # 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()
        
        # 停止 ComfyUI 进程
        await self._stop_comfyui()
        
        # 清理 socket 文件
        if _MODEL_SERVER_SOCKET and os.path.exists(_MODEL_SERVER_SOCKET):
            os.remove(_MODEL_SERVER_SOCKET)

    async def _start_comfyui(self) -> None:
        """启动 ComfyUI 子进程"""
        uvicorn_logger.info(f"Starting ComfyUI: {self._start_cmd}")
        uvicorn_logger.info(f"Work directory: {self._work_dir}")
        
        # 准备环境变量
        env = os.environ.copy()
        env.update(self._extra_env)
        
        # 构建启动命令
        cmd = self._start_cmd
        if f"--port" not in cmd:
            cmd += f" --port {self._comfyui_port}"
        if f"--listen" not in cmd:
            cmd += f" --listen {self._comfyui_host}"
        
        # 启动子进程
        self._comfyui_process = subprocess.Popen(
            cmd,
            shell=True,
            cwd=self._work_dir,
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            preexec_fn=os.setsid,  # 创建新的进程组，方便后续 kill
        )
        
        # 启动日志转发线程
        asyncio.create_task(self._forward_comfyui_logs())
        
        # 等待 ComfyUI 启动完成
        await self._wait_for_comfyui_ready()

    async def _forward_comfyui_logs(self) -> None:
        """转发 ComfyUI 的日志输出"""
        if not self._comfyui_process or not self._comfyui_process.stdout:
            return
        
        loop = asyncio.get_event_loop()
        
        def read_line():
            if self._comfyui_process and self._comfyui_process.stdout:
                return self._comfyui_process.stdout.readline()
            return b""
        
        while self._comfyui_process and self._comfyui_process.poll() is None:
            try:
                line = await loop.run_in_executor(None, read_line)
                if line:
                    uvicorn_logger.info(f"[ComfyUI] {line.decode('utf-8', errors='ignore').strip()}")
            except Exception:
                break

    async def _wait_for_comfyui_ready(self) -> None:
        """等待 ComfyUI 服务就绪"""
        uvicorn_logger.info("Waiting for ComfyUI to be ready...")
        
        start_time = time.time()
        check_interval = 1.0
        
        while time.time() - start_time < _COMFYUI_STARTUP_TIMEOUT:
            try:
                async with aiohttp.ClientSession() as session:
                    async with session.get(
                        f"{self._comfyui_url}/system_stats",
                        timeout=aiohttp.ClientTimeout(total=5)
                    ) as resp:
                        if resp.status == 200:
                            self._comfyui_ready = True
                            uvicorn_logger.info("ComfyUI is ready!")
                            return
            except Exception:
                pass
            
            # 检查进程是否还活着
            if self._comfyui_process and self._comfyui_process.poll() is not None:
                raise RuntimeError(
                    f"ComfyUI process died with exit code {self._comfyui_process.returncode}"
                )
            
            await asyncio.sleep(check_interval)
        
        raise TimeoutError(f"ComfyUI failed to start within {_COMFYUI_STARTUP_TIMEOUT}s")

    async def _stop_comfyui(self) -> None:
        """优雅停止 ComfyUI 进程"""
        if not self._comfyui_process:
            return
        
        uvicorn_logger.info("Stopping ComfyUI process...")
        
        try:
            # 先尝试 SIGTERM
            os.killpg(os.getpgid(self._comfyui_process.pid), signal.SIGTERM)
            
            # 等待进程退出
            for _ in range(10):
                if self._comfyui_process.poll() is not None:
                    uvicorn_logger.info("ComfyUI stopped gracefully")
                    return
                await asyncio.sleep(0.5)
            
            # 强制 SIGKILL
            uvicorn_logger.warning("Force killing ComfyUI...")
            os.killpg(os.getpgid(self._comfyui_process.pid), signal.SIGKILL)
            self._comfyui_process.wait()
            
        except ProcessLookupError:
            pass  # 进程已经退出
        except Exception as e:
            uvicorn_logger.error(f"Error stopping ComfyUI: {e}")

    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 路由 - 与 ComfyUI 原生 API 路径保持一致
        
        ComfyUI 原生 API:
        - GET  /                 - 前端页面
        - GET  /system_stats     - 系统状态
        - GET  /queue            - 队列信息
        - POST /prompt           - 提交工作流
        - POST /interrupt        - 中断生成
        - GET  /history          - 历史记录
        - GET  /history/{prompt_id} - 指定历史
        - GET  /view             - 查看图像
        - POST /upload/image     - 上传图像
        - GET  /embeddings       - embeddings列表
        - GET  /extensions       - 扩展列表
        - GET  /object_info      - 节点信息
        - WS   /ws               - WebSocket
        """
        
        # ========== 健康检查（额外添加）==========
        @self._app.get("/health")
        async def health():
            """健康检查接口"""
            if not self._comfyui_ready:
                return JSONResponse({"status": "initializing"}, status_code=503)
            
            if self._comfyui_process and self._comfyui_process.poll() is not None:
                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):
            """
            提交工作流 - 核心接口
            支持两种模式:
            1. 直接代理到 ComfyUI（默认）
            2. 带进度流式返回（通过 X-Stream-Response: true 头开启）
            """
            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):
            """获取指定 prompt 的历史"""
            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):
            """
            流式提交工作流 - 扩展接口
            返回 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("/prompt/sync")
        async def prompt_sync(request: Request):
            """
            同步提交工作流 - 扩展接口
            等待生成完成后返回所有图像（base64）
            """
            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({"error": data})
                    break
                if code == SERVER_CODE_STOP:
                    break
                results.append(data)
            
            return {"code": last_code, "results": results}

        # ========== 通配符路由 - 代理其他未明确定义的路径 ==========
        @self._app.api_route(
            "/{path:path}",
            methods=["GET", "POST", "PUT", "DELETE"]
        )
        async def proxy_fallback(request: Request, path: str):
            """透传其他未定义的路由到 ComfyUI"""
            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 _sse_generator(self, request: bytes) -> AsyncGenerator[str, None]:
        """SSE 响应生成器"""
        async for item in self._generate(request):
            code = item["code"]
            data = item["data"]
            yield f"{json.dumps({'code': code, 'data': data}, 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 = str(headers.get("x-trace-id", "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": "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=[]
            )
            
            # WebSocket 监听进度
            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_msg}
                            return
                    
                    elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
                        yield {"code": SERVER_CODE_ERROR, "data": "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,  # base64 编码的图像列表
                    "prompt_id": prompt_id
                }
            }
            yield {"code": SERVER_CODE_STOP, "data": ""}
            
        except asyncio.TimeoutError:
            logger.error("Generation timeout")
            yield {"code": SERVER_CODE_ERROR, "data": "timeout"}
        except Exception as e:
            logger.error(f"Error: {e}")
            yield {"code": SERVER_CODE_ERROR, "data": 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]:
        """提交工作流到 ComfyUI"""
        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]:
        """获取生成的图像（base64）"""
        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