# /root/code/python/ieops-python/sdk/src/handler/server/diffuser.py
"""
Diffuser Server - Hugging Face Diffusers 推理服务接口层

设计模式与 OpenAIServer 一致：
- Server 只负责 HTTP 接口层
- 推理引擎初始化和推理逻辑由 handler 中的 Model 类实现

架构：
┌─────────────────────────────────────────────────────────────┐
│  DiffuserServer (SDK - 通用接口层)                           │
│  - HTTP 接口 (/v1/txt2img, /v1/img2img, /infer 等)          │
│  - 请求解析、响应格式化                                      │
│  - 接收 txt2img_func, img2img_func 等推理函数                │
└─────────────────────────────────────────────────────────────┘
                              ↓ 调用 func
┌─────────────────────────────────────────────────────────────┐
│  DiffuserModel 类 (Worker handler.py)                       │
│  - Pipeline 初始化                                          │
│  - 推理逻辑 (async generator)                               │
└─────────────────────────────────────────────────────────────┘
"""

import asyncio
import json
import os
import time
from typing import Any, AsyncGenerator, Callable, Dict, List, Optional
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, Field

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 DiffuserError(Exception):
    """Diffuser 推理错误"""
    def __init__(self, message: str):
        super().__init__(message)
        self.message = message


# ========== Pydantic 模型定义 ==========
class Txt2ImgRequest(BaseModel):
    """文生图请求"""
    prompt: str = Field(..., description="正向提示词")
    negative_prompt: str = Field(default="", description="负向提示词")
    width: int = Field(default=1024, ge=256, le=4096, description="图像宽度")
    height: int = Field(default=1024, ge=256, le=4096, description="图像高度")
    num_inference_steps: int = Field(default=30, ge=1, le=150, description="推理步数")
    guidance_scale: float = Field(default=7.5, ge=0, le=30, description="引导系数")
    num_images: int = Field(default=1, ge=1, le=4, description="生成图片数量")
    seed: Optional[int] = Field(default=None, description="随机种子")
    scheduler: Optional[str] = Field(default=None, description="调度器类型")
    lora_weights: Optional[List[Dict[str, Any]]] = Field(default=None, description="LoRA 权重配置")


class Img2ImgRequest(BaseModel):
    """图生图请求"""
    prompt: str = Field(..., description="正向提示词")
    negative_prompt: str = Field(default="", description="负向提示词")
    image: str = Field(..., description="输入图像 (base64)")
    strength: float = Field(default=0.75, ge=0, le=1, description="去噪强度")
    num_inference_steps: int = Field(default=30, ge=1, le=150, description="推理步数")
    guidance_scale: float = Field(default=7.5, ge=0, le=30, description="引导系数")
    num_images: int = Field(default=1, ge=1, le=4, description="生成图片数量")
    seed: Optional[int] = Field(default=None, description="随机种子")


class InpaintRequest(BaseModel):
    """图像修复请求"""
    prompt: str = Field(..., description="正向提示词")
    negative_prompt: str = Field(default="", description="负向提示词")
    image: str = Field(..., description="输入图像 (base64)")
    mask_image: str = Field(..., description="遮罩图像 (base64)")
    width: Optional[int] = Field(default=None, description="输出宽度")
    height: Optional[int] = Field(default=None, description="输出高度")
    num_inference_steps: int = Field(default=30, ge=1, le=150, description="推理步数")
    guidance_scale: float = Field(default=7.5, ge=0, le=30, description="引导系数")
    strength: float = Field(default=0.75, ge=0, le=1, description="去噪强度")
    num_images: int = Field(default=1, ge=1, le=4, description="生成图片数量")
    seed: Optional[int] = Field(default=None, description="随机种子")


# 推理函数类型定义
Txt2ImgFunc = Callable[[Dict[str, Any]], AsyncGenerator[Dict[str, Any], None]]
Img2ImgFunc = Callable[[Dict[str, Any]], AsyncGenerator[Dict[str, Any], None]]
InpaintFunc = Callable[[Dict[str, Any]], AsyncGenerator[Dict[str, Any], None]]
ModelInfoFunc = Callable[[], Dict[str, Any]]


class DiffuserServer(BaseServer):
    """
    Diffuser 推理服务器 - 只负责接口层
    
    与 OpenAIServer 设计模式一致：
    - 接收推理函数作为参数
    - 只负责 HTTP 接口、请求解析、响应格式化
    - 推理引擎初始化和推理逻辑由 handler 实现
    
    使用示例:
    ```python
    # handler.py
    class DiffuserModel:
        def __init__(self):
            self.pipe = AutoPipelineForText2Image.from_pretrained(...)
        
        async def txt2img(self, query: dict):
            # 推理逻辑，yield 进度和结果
            yield {"type": "progress", "progress": 0.5, "step": 15, "total_steps": 30}
            yield {"type": "result", "images": [...], "seed": 12345}
    
    model = DiffuserModel()
    server = DiffuserServer(
        txt2img_func=model.txt2img,
        img2img_func=model.img2img,
        model_info_func=model.get_info,
    )
    Handler(server=server).serve()
    ```
    """
    
    def __init__(
        self,
        txt2img_func: Optional[Txt2ImgFunc] = None,
        img2img_func: Optional[Img2ImgFunc] = None,
        inpaint_func: Optional[InpaintFunc] = None,
        model_info_func: Optional[ModelInfoFunc] = None,
    ) -> None:
        """
        Args:
            txt2img_func: 文生图推理函数 (async generator)
            img2img_func: 图生图推理函数 (async generator)
            inpaint_func: 图像修复推理函数 (async generator)
            model_info_func: 获取模型信息函数
        """
        self._txt2img_func = txt2img_func
        self._img2img_func = img2img_func
        self._inpaint_func = inpaint_func
        self._model_info_func = model_info_func
        
        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("Diffuser Server starting...")
        yield
        uvicorn_logger.info("Diffuser Server shutting down...")
        
        # 清理 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():
            """健康检查"""
            return {"status": "ok"}

        @self._app.get("/")
        async def root():
            """根路径"""
            return {
                "service": "Diffuser Server",
                "status": "ready",
                "endpoints": {
                    "health": "/health",
                    "txt2img": "/v1/txt2img",
                    "img2img": "/v1/img2img",
                    "inpaint": "/v1/inpaint",
                    "model_info": "/v1/model/info",
                },
                "features": {
                    "txt2img": self._txt2img_func is not None,
                    "img2img": self._img2img_func is not None,
                    "inpaint": self._inpaint_func is not None,
                }
            }

        @self._app.get("/v1/model/info")
        async def get_model_info():
            """获取模型信息"""
            if self._model_info_func:
                return self._model_info_func()
            return {"message": "model_info_func not provided"}

        # ========== 文生图 ==========
        @self._app.post("/v1/txt2img")
        async def txt2img(request: Txt2ImgRequest):
            """文生图（同步）"""
            if not self._txt2img_func:
                raise HTTPException(status_code=501, detail="txt2img not implemented")
            
            result = None
            query = request.model_dump()
            query["trace_id"] = f"traceid-{await arandstr(8)}"
            
            async for data in self._txt2img_func(query):
                if isinstance(data, DiffuserError):
                    raise HTTPException(status_code=500, detail=data.message)
                if data.get("type") == "result":
                    result = data
                elif data.get("type") == "error":
                    raise HTTPException(status_code=500, detail=data.get("error"))
            return result

        @self._app.post("/v1/txt2img/stream")
        async def txt2img_stream(request: Txt2ImgRequest):
            """文生图（流式）"""
            if not self._txt2img_func:
                raise HTTPException(status_code=501, detail="txt2img not implemented")
            
            query = request.model_dump()
            query["trace_id"] = f"traceid-{await arandstr(8)}"
            
            return StreamingResponse(
                self._sse_generator(self._txt2img_func, query),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )

        # ========== 图生图 ==========
        @self._app.post("/v1/img2img")
        async def img2img(request: Img2ImgRequest):
            """图生图（同步）"""
            if not self._img2img_func:
                raise HTTPException(status_code=501, detail="img2img not implemented")
            
            result = None
            query = request.model_dump()
            query["trace_id"] = f"traceid-{await arandstr(8)}"
            
            async for data in self._img2img_func(query):
                if isinstance(data, DiffuserError):
                    raise HTTPException(status_code=500, detail=data.message)
                if data.get("type") == "result":
                    result = data
                elif data.get("type") == "error":
                    raise HTTPException(status_code=500, detail=data.get("error"))
            return result

        @self._app.post("/v1/img2img/stream")
        async def img2img_stream(request: Img2ImgRequest):
            """图生图（流式）"""
            if not self._img2img_func:
                raise HTTPException(status_code=501, detail="img2img not implemented")
            
            query = request.model_dump()
            query["trace_id"] = f"traceid-{await arandstr(8)}"
            
            return StreamingResponse(
                self._sse_generator(self._img2img_func, query),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )

        # ========== 图像修复 ==========
        @self._app.post("/v1/inpaint")
        async def inpaint(request: InpaintRequest):
            """图像修复（同步）"""
            if not self._inpaint_func:
                raise HTTPException(status_code=501, detail="inpaint not implemented")
            
            result = None
            query = request.model_dump()
            query["trace_id"] = f"traceid-{await arandstr(8)}"
            
            async for data in self._inpaint_func(query):
                if isinstance(data, DiffuserError):
                    raise HTTPException(status_code=500, detail=data.message)
                if data.get("type") == "result":
                    result = data
                elif data.get("type") == "error":
                    raise HTTPException(status_code=500, detail=data.get("error"))
            return result

        @self._app.post("/v1/inpaint/stream")
        async def inpaint_stream(request: InpaintRequest):
            """图像修复（流式）"""
            if not self._inpaint_func:
                raise HTTPException(status_code=501, detail="inpaint not implemented")
            
            query = request.model_dump()
            query["trace_id"] = f"traceid-{await arandstr(8)}"
            
            return StreamingResponse(
                self._sse_generator(self._inpaint_func, query),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )

        # ========== 兼容 ieops 协议 ==========
        @self._app.post("/sse/infer")
        async def infer_sse(request: Request):
            """SSE 推理接口 - 兼容 ieops 协议"""
            body = await request.body()
            return StreamingResponse(
                self._ieops_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):
            """同步推理接口 - 兼容 ieops 协议"""
            body = await request.body()
            results: List[Any] = []
            last_code = SERVER_CODE_OK
            
            async for item in self._ieops_generator(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}

    async def _sse_generator(
        self, 
        func: Callable[[Dict[str, Any]], AsyncGenerator[Dict[str, Any], None]], 
        query: Dict[str, Any]
    ) -> AsyncGenerator[str, None]:
        """SSE 响应生成器"""
        try:
            async for data in func(query):
                if isinstance(data, DiffuserError):
                    yield f"data: {json.dumps({'type': 'error', 'error': data.message})}\n\n"
                    return
                yield f"data: {json.dumps(data)}\n\n"
        except Exception as e:
            uvicorn_logger.error(f"SSE generator error: {e}")
            yield f"data: {json.dumps({'type': 'error', 'error': str(e)})}\n\n"

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

    async def _ieops_generator(self, request: bytes) -> AsyncGenerator[Dict[str, Any], None]:
        """ieops 协议生成器"""
        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("Diffuser request received")
        
        # 解析任务类型
        if isinstance(payload, dict):
            task_type = payload.get("type", "txt2img")
            query = payload.get("params", payload)
        else:
            query = json.loads(payload) if isinstance(payload, str) else payload
            task_type = query.get("type", "txt2img")
        
        query["trace_id"] = trace_id
        
        # 选择推理函数
        func_map = {
            "txt2img": self._txt2img_func,
            "img2img": self._img2img_func,
            "inpaint": self._inpaint_func,
        }
        func = func_map.get(task_type)
        
        if not func:
            yield {"code": SERVER_CODE_ERROR, "data": {"error": f"Unknown or unsupported task type: {task_type}"}}
            return
        
        try:
            async for data in func(query):
                if isinstance(data, DiffuserError):
                    yield {"code": SERVER_CODE_ERROR, "data": {"error": data.message}}
                    return
                elif data.get("type") == "error":
                    yield {"code": SERVER_CODE_ERROR, "data": {"error": data.get("error", "Unknown error")}}
                    return
                else:
                    # 直接返回 handler 的数据结构，不做额外包装
                    yield {"code": SERVER_CODE_OK, "data": data}
                    if data.get("type") == "result":
                        yield {"code": SERVER_CODE_STOP, "data": None}
                        return
            
            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")
