"""
OpenAI Server - OpenAI 兼容 API 接口层

提供与 OpenAI 官方 API 一致的接口：
- /v1/chat/completions - Chat Completions API
- /v1/completions - Completions API
- /v1/models - List models

同时保留 ieops 协议兼容接口：
- /sse/infer - SSE 流式推理
- /infer - 同步推理
"""
from typing import TypedDict, Callable, Any, Generator, Optional, Dict, List, AsyncGenerator, Union
from concurrent.futures import ThreadPoolExecutor
import asyncio
import uvicorn
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from contextlib import asynccontextmanager
from pydantic import BaseModel, Field
from .server import BaseServer
from ...utils.log import uvicorn_logger, log
from ...utils.util import APP_ID, arandstr
import inspect
import threading
import traceback
import time
import json
import queue
import os
from ...handler.tokenizer import TokenizerPool


_APP_NAME = os.getenv("APP_NAME", "ieops")
_MODEL_NAME = os.getenv("MODEL_NAME", "gpt-3.5-turbo")

_MODEL_SERVER_SOCKET = os.getenv("MODEL_SERVER_SOCKET", f"{APP_ID}.sock")
_MODEL_SERVER_PORT = os.getenv("MODEL_SERVER_PORT", 8001)
_MODEL_SERVER_HOST = os.getenv("MODEL_SERVER_HOST", "127.0.0.1")

_THREAD_CONCURRENCY = os.getenv("MODEL_THREAD_CONCURRENCY", "8")
_MODEL_TIMEOUT = os.getenv("MODEL_TIMEOUT", 600)
_GRACEFUL_SHUTDOWN_TIME = os.getenv("GRACEFUL_SHUTDOWN_TIME", 3)

ENGINE_READY: bool = False

# ieops 协议状态码
SERVER_CODE_OK = 0
SERVER_CODE_ERROR = 1
SERVER_CODE_STOP = 2


class ModelError(Exception):
    """模型错误"""
    def __init__(self, message: str):
        super().__init__(message)
        self.message = message


# ========== OpenAI 兼容的 Pydantic 模型 ==========

class ChatMessage(BaseModel):
    """聊天消息"""
    role: str = Field(..., description="消息角色: system, user, assistant")
    content: str = Field(..., description="消息内容")
    name: Optional[str] = Field(default=None, description="发送者名称")


class ChatCompletionRequest(BaseModel):
    """Chat Completions 请求"""
    model: str = Field(default=_MODEL_NAME, description="模型名称")
    messages: List[ChatMessage] = Field(..., description="消息列表")
    temperature: Optional[float] = Field(default=1.0, ge=0, le=2, description="采样温度")
    top_p: Optional[float] = Field(default=1.0, ge=0, le=1, description="核采样")
    n: Optional[int] = Field(default=1, ge=1, description="生成数量")
    stream: Optional[bool] = Field(default=False, description="是否流式返回")
    stop: Optional[Union[str, List[str]]] = Field(default=None, description="停止词")
    max_tokens: Optional[int] = Field(default=None, description="最大 token 数")
    presence_penalty: Optional[float] = Field(default=0, ge=-2, le=2)
    frequency_penalty: Optional[float] = Field(default=0, ge=-2, le=2)
    user: Optional[str] = Field(default=None, description="用户标识")


class CompletionRequest(BaseModel):
    """Completions 请求"""
    model: str = Field(default=_MODEL_NAME, description="模型名称")
    prompt: Union[str, List[str]] = Field(..., description="提示词")
    temperature: Optional[float] = Field(default=1.0, ge=0, le=2)
    top_p: Optional[float] = Field(default=1.0, ge=0, le=1)
    n: Optional[int] = Field(default=1, ge=1)
    stream: Optional[bool] = Field(default=False)
    stop: Optional[Union[str, List[str]]] = Field(default=None)
    max_tokens: Optional[int] = Field(default=16)
    presence_penalty: Optional[float] = Field(default=0, ge=-2, le=2)
    frequency_penalty: Optional[float] = Field(default=0, ge=-2, le=2)
    user: Optional[str] = Field(default=None)


class QueueItem(TypedDict):
    code: int
    data: Any


class OpenAIServer(BaseServer):
    """
    OpenAI 兼容 API 服务器
    
    提供与 OpenAI 官方一致的接口，同时保留 ieops 协议兼容。
    
    使用示例:
    ```python
    class Model:
        async def handler(self, query):
            # query 包含: messages, temperature, max_tokens 等
            yield {
                "results": [{"content": "Hello", "role": "assistant"}],
                "stop_reasons": [""],
                "stoped": False,
            }
    
    model = Model()
    server = OpenAIServer(func=model.handler, model_name="my-model")
    Handler(server=server).serve()
    ```
    """
    
    def __init__(
        self, 
        func: Callable[[Any], Generator[Any, bool, None]], 
        abort: Optional[Callable[[str], None]] = None,
        model_name: Optional[str] = None,
        model_info: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Args:
            func: 推理函数 (async generator)
            abort: 中断函数
            model_name: 模型名称
            model_info: 模型信息（用于 /v1/models 接口）
        """
        self._func = func
        self._cancelled_events: Dict[str, bool] = {}
        self._abort = abort
        self._model_name = model_name or _MODEL_NAME
        self._model_info = model_info or {}
        self._executor = ThreadPoolExecutor(max_workers=int(_THREAD_CONCURRENCY) * 2 + 3)
        self._app = FastAPI(lifespan=self._lifespan)
        self._server: Optional[uvicorn.Server] = None
        self._setup_routes()

    @asynccontextmanager
    async def _lifespan(self, app: FastAPI):
        uvicorn_logger.info("OpenAI Server starting...")
        yield
        uvicorn_logger.info("OpenAI Server shutting down...")
        for trace_id in list(self._cancelled_events.keys()):
            self._cancelled_events[trace_id] = True
        await asyncio.sleep(0.5)
        self._executor.shutdown(wait=True)
        try:
            TokenizerPool(None).shutdown(wait=True)
        except Exception as e:
            uvicorn_logger.error(f"Error shutting down TokenizerPool: {e}")
        finally:
            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=int(_GRACEFUL_SHUTDOWN_TIME)
            )
        else:
            config = uvicorn.Config(
                self._app,
                host=_MODEL_SERVER_HOST,
                port=int(_MODEL_SERVER_PORT),
                log_level="info",
                loop="asyncio",
                timeout_graceful_shutdown=int(_GRACEFUL_SHUTDOWN_TIME)
            )
        
        self._server = uvicorn.Server(config)
        await self._server.serve()

    def _setup_routes(self) -> None:
        """设置 API 路由"""
        
        # ========== 基础接口 ==========
        @self._app.get("/")
        async def root():
            return {
                "message": "OpenAI Compatible Server",
                "status": "ready" if ENGINE_READY else "initializing",
                "model": self._model_name,
                "endpoints": {
                    "health": "/health",
                    "models": "/v1/models",
                    "chat_completions": "/v1/chat/completions",
                    "completions": "/v1/completions",
                }
            }

        @self._app.get("/health")
        async def health():
            return {"status": "ok"}

        # ========== OpenAI 兼容接口 ==========
        
        @self._app.get("/v1/models")
        async def list_models():
            """列出可用模型"""
            return {
                "object": "list",
                "data": [{
                    "id": self._model_name,
                    "object": "model",
                    "created": int(time.time()),
                    "owned_by": "organization",
                    **self._model_info
                }]
            }

        @self._app.get("/v1/models/{model}")
        async def retrieve_model(model: str):
            """获取模型信息"""
            if model != self._model_name:
                raise HTTPException(status_code=404, detail=f"Model '{model}' not found")
            return {
                "id": self._model_name,
                "object": "model",
                "created": int(time.time()),
                "owned_by": "organization",
                **self._model_info
            }

        @self._app.post("/v1/chat/completions")
        async def chat_completions(request: ChatCompletionRequest):
            """Chat Completions API - OpenAI 兼容"""
            request_id = f"chatcmpl-{await arandstr(24)}"
            created = int(time.time())
            
            # 转换为内部格式
            query = {
                "messages": [msg.model_dump() for msg in request.messages],
                "model": request.model,
                "temperature": request.temperature,
                "top_p": request.top_p,
                "max_tokens": request.max_tokens,
                "stop": request.stop,
                "n": request.n,
                "stream": request.stream,
                "trace_id": f"traceid-{await arandstr(8)}",
            }
            
            if request.stream:
                return StreamingResponse(
                    self._openai_stream_generator(query, request_id, created, "chat.completion.chunk"),
                    media_type="text/event-stream",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                        "X-Accel-Buffering": "no",
                    }
                )
            else:
                return await self._openai_chat_response(query, request_id, created)

        @self._app.post("/v1/completions")
        async def completions(request: CompletionRequest):
            """Completions API - OpenAI 兼容"""
            request_id = f"cmpl-{await arandstr(24)}"
            created = int(time.time())
            
            # 将 prompt 转换为 messages 格式
            prompt = request.prompt if isinstance(request.prompt, str) else request.prompt[0]
            query = {
                "messages": [{"role": "user", "content": prompt}],
                "model": request.model,
                "temperature": request.temperature,
                "top_p": request.top_p,
                "max_tokens": request.max_tokens,
                "stop": request.stop,
                "n": request.n,
                "stream": request.stream,
                "trace_id": f"traceid-{await arandstr(8)}",
                "is_completion": True,  # 标记为 completion 模式
            }
            
            if request.stream:
                return StreamingResponse(
                    self._openai_stream_generator(query, request_id, created, "text_completion"),
                    media_type="text/event-stream",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                        "X-Accel-Buffering": "no",
                    }
                )
            else:
                return await self._openai_completion_response(query, request_id, created)

    # ========== OpenAI 格式响应生成 ==========
    
    async def _openai_chat_response(
        self, 
        query: Dict[str, Any], 
        request_id: str, 
        created: int
    ) -> Dict[str, Any]:
        """生成 OpenAI Chat Completions 格式的响应"""
        content_parts = []
        finish_reason = "stop"
        prompt_tokens = 0
        completion_tokens = 0
        
        try:
            async for v in self._run_inference(query):
                if isinstance(v, ModelError):
                    raise HTTPException(status_code=500, detail=v.message)
                if v is None:
                    break
                
                # 提取内容
                if "results" in v:
                    for ret in v["results"]:
                        if ret.get("content"):
                            content_parts.append(ret["content"])
                
                # 提取 token 统计
                if "token_encode" in v:
                    prompt_tokens = v["token_encode"]
                if "token_decode" in v:
                    completion_tokens = v["token_decode"]
                
                # 提取停止原因
                if "stop_reasons" in v and v["stop_reasons"]:
                    reason = v["stop_reasons"][0]
                    if reason and reason != "":
                        finish_reason = reason if reason in ("stop", "length", "content_filter") else "stop"
        
        except asyncio.TimeoutError:
            raise HTTPException(status_code=504, detail="Request timeout")
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))
        
        return {
            "id": request_id,
            "object": "chat.completion",
            "created": created,
            "model": self._model_name,
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": "".join(content_parts)
                },
                "finish_reason": finish_reason
            }],
            "usage": {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens
            }
        }

    async def _openai_completion_response(
        self, 
        query: Dict[str, Any], 
        request_id: str, 
        created: int
    ) -> Dict[str, Any]:
        """生成 OpenAI Completions 格式的响应"""
        content_parts = []
        finish_reason = "stop"
        prompt_tokens = 0
        completion_tokens = 0
        
        try:
            async for v in self._run_inference(query):
                if isinstance(v, ModelError):
                    raise HTTPException(status_code=500, detail=v.message)
                if v is None:
                    break
                
                if "results" in v:
                    for ret in v["results"]:
                        if ret.get("content"):
                            content_parts.append(ret["content"])
                
                if "token_encode" in v:
                    prompt_tokens = v["token_encode"]
                if "token_decode" in v:
                    completion_tokens = v["token_decode"]
                
                if "stop_reasons" in v and v["stop_reasons"]:
                    reason = v["stop_reasons"][0]
                    if reason and reason != "":
                        finish_reason = reason if reason in ("stop", "length") else "stop"
        
        except asyncio.TimeoutError:
            raise HTTPException(status_code=504, detail="Request timeout")
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))
        
        return {
            "id": request_id,
            "object": "text_completion",
            "created": created,
            "model": self._model_name,
            "choices": [{
                "index": 0,
                "text": "".join(content_parts),
                "finish_reason": finish_reason
            }],
            "usage": {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens
            }
        }

    async def _openai_stream_generator(
        self, 
        query: Dict[str, Any], 
        request_id: str, 
        created: int,
        object_type: str
    ) -> AsyncGenerator[str, None]:
        """生成 OpenAI SSE 流式响应"""
        is_chat = object_type == "chat.completion.chunk"
        
        try:
            async for v in self._run_inference(query):
                if isinstance(v, ModelError):
                    error_data = {
                        "error": {
                            "message": v.message,
                            "type": "model_error",
                            "code": "model_error"
                        }
                    }
                    yield f"data: {json.dumps(error_data)}\n\n"
                    return
                
                if v is None:
                    break
                
                # 提取内容增量
                content = ""
                if "results" in v:
                    for ret in v["results"]:
                        if ret.get("content"):
                            content = ret["content"]
                
                # 检查是否结束
                finish_reason = None
                if v.get("stoped") or (v.get("stop_reasons") and v["stop_reasons"][0]):
                    reason = v.get("stop_reasons", ["stop"])[0]
                    finish_reason = reason if reason in ("stop", "length", "content_filter") else "stop"
                
                # 构建 chunk
                if is_chat:
                    chunk = {
                        "id": request_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": self._model_name,
                        "choices": [{
                            "index": 0,
                            "delta": {"content": content} if content else {},
                            "finish_reason": finish_reason
                        }]
                    }
                else:
                    chunk = {
                        "id": request_id,
                        "object": "text_completion",
                        "created": created,
                        "model": self._model_name,
                        "choices": [{
                            "index": 0,
                            "text": content,
                            "finish_reason": finish_reason
                        }]
                    }
                
                yield f"data: {json.dumps(chunk)}\n\n"
                
                if finish_reason:
                    break
            
            yield "data: [DONE]\n\n"
            
        except asyncio.TimeoutError:
            error_data = {"error": {"message": "Request timeout", "type": "timeout", "code": "timeout"}}
            yield f"data: {json.dumps(error_data)}\n\n"
        except Exception as e:
            error_data = {"error": {"message": str(e), "type": "server_error", "code": "server_error"}}
            yield f"data: {json.dumps(error_data)}\n\n"

    async def _run_inference(self, query: Dict[str, Any]) -> AsyncGenerator[Any, None]:
        """运行推理"""
        trace_id = query.get("trace_id", f"traceid-{await arandstr(8)}")
        query["trace_id"] = trace_id
        
        if trace_id in self._cancelled_events:
            return
        
        self._cancelled_events[trace_id] = False
        
        try:
            if inspect.isasyncgenfunction(self._func):
                threading.Thread(target=self._time_counter, args=(trace_id, int(_MODEL_TIMEOUT))).start()
                async for v in self._func(query):
                    if self._cancelled_events.get(trace_id):
                        if self._abort is not None:
                            await self._abort(trace_id)
                        raise asyncio.TimeoutError
                    yield v
            else:
                loop = asyncio.get_event_loop()
                q: "queue.Queue[Dict[str, Any]]" = queue.Queue()
                await asyncio.wait_for(
                    loop.run_in_executor(self._executor, self._func_caller, trace_id, query, q),
                    int(_MODEL_TIMEOUT) // 10
                )
                while True:
                    v = await loop.run_in_executor(self._executor, q.get)
                    if v.get("code") == SERVER_CODE_STOP:
                        break
                    if v.get("code") == SERVER_CODE_ERROR:
                        raise Exception(v.get("data", {}).get("error", "Unknown error"))
                    yield v.get("data", v)
        finally:
            if trace_id in self._cancelled_events:
                del self._cancelled_events[trace_id]

    def _func_caller(self, trace_id: str, payload: Dict[str, Any], q: "queue.Queue[Dict[str, Any]]"):
        """同步函数调用器"""
        try:
            it = self._func(payload)
            if not hasattr(it, "send"):
                q.put({"code": SERVER_CODE_OK, "data": it})
                q.put({"code": SERVER_CODE_STOP, "data": None})
                return
            v = it.send(None)
            while True:
                q.put({"code": SERVER_CODE_OK, "data": v})
                if self._cancelled_events.get(trace_id):
                    try:
                        v = it.send(True)
                    except StopIteration:
                        break
                    break
                else:
                    try:
                        v = it.send(False)
                    except StopIteration:
                        break
        except Exception:
            log.get_logger(trace_id=trace_id).error("error: {}", traceback.format_exc())
            q.put({"code": SERVER_CODE_ERROR, "data": {"error": traceback.format_exc()}})
        else:
            q.put({"code": SERVER_CODE_STOP, "data": None})
        finally:
            log.get_logger(trace_id=trace_id).info("func_caller finished")
            if trace_id in self._cancelled_events:
                del self._cancelled_events[trace_id]

    def _time_counter(self, trace_id: str, timeout: int):
        """超时计数器"""
        for _ in range(timeout):
            if self._cancelled_events.get(trace_id):
                log.get_logger(trace_id=trace_id).info("time counter stop")
                if trace_id in self._cancelled_events:
                    del self._cancelled_events[trace_id]
                return
            time.sleep(0.1)
        if trace_id in self._cancelled_events:
            self._cancelled_events[trace_id] = True
