#! /usr/bin/env python3
# -*- coding: utf-8 -*-

"""
LLM Server - 兼容 OpenAI 的 Chat Completions API 服务

基于 FastAPI 的 HTTP 服务封装，使用 LLMBackend 进行推理。
"""

import time
import uuid
from contextlib import asynccontextmanager
from typing import List, Literal, Optional, Union

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from loguru import logger
from pydantic import BaseModel, Field

from .base import ChatMessage, GenerateConfig, ModelConfig, BaseLLMBackend
from .batch import BatchScheduler
from .backend import TransformersBackend


# ============== API 响应模型 ==============


class ChatCompletionChoice(BaseModel):
    """响应选项"""
    index: int = 0
    message: ChatMessage
    finish_reason: Optional[str] = "stop"


class UsageInfo(BaseModel):
    """Token 使用统计"""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class ChatCompletionResponse(BaseModel):
    """Chat Completion 响应"""
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:8]}")
    object: Literal["chat.completion"] = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = ""
    choices: List[ChatCompletionChoice] = Field(default_factory=list)
    usage: UsageInfo = Field(default_factory=UsageInfo)


class DeltaMessage(BaseModel):
    """流式响应的增量消息"""
    role: Optional[str] = None
    content: Optional[str] = None


class ChatCompletionChunkChoice(BaseModel):
    """流式响应选项"""
    index: int = 0
    delta: DeltaMessage
    finish_reason: Optional[str] = None


class ChatCompletionChunk(BaseModel):
    """流式响应块"""
    id: str = ""
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = ""
    choices: List[ChatCompletionChunkChoice] = Field(default_factory=list)


class ChatCompletionRequest(BaseModel):
    """Chat Completion 请求"""
    model: str = Field(..., description="模型名称")
    messages: List[ChatMessage] = Field(..., description="消息列表")
    temperature: float = Field(default=0.7, ge=0, le=2)
    top_p: float = Field(default=0.9, ge=0, le=1)
    max_tokens: Optional[int] = Field(default=512)
    stream: bool = Field(default=False)
    stop: Optional[Union[str, List[str]]] = None


class ModelInfo(BaseModel):
    """模型信息"""
    id: str
    object: Literal["model"] = "model"
    created: int = 0
    owned_by: str = "local"


class ModelsResponse(BaseModel):
    """模型列表响应"""
    object: Literal["list"] = "list"
    data: List[ModelInfo] = Field(default_factory=list)


# ============== Server ==============


class LLMServer:
    """LLM HTTP 服务

    Args:
        backend: LLM 后端实例，默认使用 TransformersBackend
        model: 预加载的模型 ID
        device: 设备类型
        local_dir: 本地模型目录
        dtype: 数据类型 (float16/bfloat16/float32)
        attn: 注意力实现 (eager/sdpa/flash_attention_2)
        model_class: 模型类名
        processor_class: 处理器类名
        vision_processor: 视觉处理器类型 (qwen_vl/general)
        chat_template_kwargs: chat template 额外参数
    """

    def __init__(
        self,
        backend: Optional[BaseLLMBackend] = None,
        model: Optional[str] = None,
        device: Optional[str] = None,
        local_dir: Optional[str] = None,
        dtype: Optional[str] = None,
        attn: Optional[str] = None,
        model_class: Optional[str] = None,
        processor_class: Optional[str] = None,
        vision_processor: Optional[str] = None,
        chat_template_kwargs: Optional[dict] = None,
        lora_modules: Optional[dict] = None,
        # 推理优化参数
        compile: bool = False,
        compile_mode: str = "reduce-overhead",
        cache: Optional[str] = None,
        # 批处理参数
        batch_max_size: int = 0,
        batch_wait_s: float = 1.0,
    ):
        self.backend = backend or TransformersBackend()
        if device:
            self.backend._device = device

        self._preload_model = model
        self._local_dir = local_dir
        self._dtype = dtype
        self._attn = attn
        self._model_class = model_class
        self._processor_class = processor_class
        self._vision_processor = vision_processor
        self._chat_template_kwargs = chat_template_kwargs or {}
        self._lora_modules = lora_modules or {}
        self._compile = compile
        self._compile_mode = compile_mode
        self._cache = cache
        # 批处理：batch_max_size > 0 时启用
        self._batch_max_size = batch_max_size
        self._batch_wait_s = batch_wait_s
        self._scheduler: Optional[BatchScheduler] = None
        self.app = self._create_app()

    def _create_app(self) -> FastAPI:
        """创建 FastAPI 应用"""

        @asynccontextmanager
        async def lifespan(app: FastAPI):
            if self._preload_model:
                try:
                    config = ModelConfig(
                        model_id=self._preload_model,
                        local_dir=self._local_dir,
                        torch_dtype=self._dtype,
                        attn_implementation=self._attn,
                        model_class=self._model_class,
                        processor_class=self._processor_class,
                        vision_processor=self._vision_processor,
                        chat_template_kwargs=self._chat_template_kwargs,
                        lora_modules=self._lora_modules,
                        compile=self._compile,
                        compile_mode=self._compile_mode,
                    )
                    await self.backend.load_model(config)
                except Exception as e:
                    logger.error(f"Failed to load {self._preload_model}: {e}")
                    raise

            # 启动批处理调度器
            if self._batch_max_size > 0:
                self._scheduler = BatchScheduler(
                    self.backend,
                    max_batch_size=self._batch_max_size,
                    max_wait_s=self._batch_wait_s,
                )
                await self._scheduler.start()

            yield

            # 停止批处理调度器
            if self._scheduler:
                await self._scheduler.stop()

        app = FastAPI(
            title="LLM Server",
            description="OpenAI Compatible Chat Completions API",
            version="1.0.0",
            lifespan=lifespan,
        )

        app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

        self._register_routes(app)
        return app

    def _register_routes(self, app: FastAPI) -> None:
        """注册路由"""

        @app.get("/health")
        async def health():
            return {"status": "ok", "model": self.backend.model_id}

        @app.get("/v1/models", response_model=ModelsResponse)
        async def list_models():
            models = []
            if self.backend.model_id:
                models.append(ModelInfo(id=self.backend.model_id, created=int(time.time())))
            # 列出 LoRA 适配器
            for name in self._lora_modules:
                models.append(ModelInfo(id=name, created=int(time.time()), owned_by="lora"))
            return ModelsResponse(data=models)

        @app.post("/v1/chat/completions")
        async def chat_completions(request: ChatCompletionRequest):
            return await self._handle_chat(request)

    async def _handle_chat(self, request: ChatCompletionRequest):
        """处理 chat 请求"""
        if not self.backend.is_loaded:
            raise HTTPException(status_code=503, detail="Model not loaded")

        # LoRA 路由：请求中的 model 字段匹配 LoRA 适配器名称
        if self._lora_modules and hasattr(self.backend, "set_active_lora"):
            if request.model in self._lora_modules:
                self.backend.set_active_lora(request.model)
            else:
                self.backend.set_active_lora(None)

        stop = request.stop if isinstance(request.stop, list) else (
            [request.stop] if request.stop else None
        )

        gen_config = GenerateConfig(
            max_tokens=request.max_tokens or 512,
            temperature=request.temperature,
            top_p=request.top_p,
            stop=stop,
            cache_implementation=self._cache,
        )

        try:
            if request.stream:
                return await self._stream_response(request, gen_config)
            else:
                return await self._normal_response(request, gen_config)
        except Exception as e:
            logger.exception(f"Chat error: {e}")
            raise HTTPException(status_code=500, detail=str(e))

    async def _normal_response(
        self, request: ChatCompletionRequest, config: GenerateConfig
    ) -> ChatCompletionResponse:
        """普通响应"""
        if self._scheduler:
            text, prompt_tokens, completion_tokens = await self._scheduler.submit(
                messages=request.messages,
                config=config,
            )
        else:
            text, prompt_tokens, completion_tokens = await self.backend.generate(
                messages=request.messages,
                config=config,
            )

        return ChatCompletionResponse(
            model=request.model,
            choices=[
                ChatCompletionChoice(
                    message=ChatMessage(role="assistant", content=text),
                    finish_reason="stop",
                )
            ],
            usage=UsageInfo(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=prompt_tokens + completion_tokens,
            ),
        )

    async def _stream_response(
        self, request: ChatCompletionRequest, config: GenerateConfig
    ) -> StreamingResponse:
        """流式响应"""
        response_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"

        async def generate():
            # 发送角色
            chunk = ChatCompletionChunk(
                id=response_id,
                model=request.model,
                choices=[
                    ChatCompletionChunkChoice(
                        delta=DeltaMessage(role="assistant"),
                    )
                ],
            )
            yield f"data: {chunk.model_dump_json()}\n\n"

            # 流式内容
            async for text in self.backend.generate_stream(
                messages=request.messages,
                config=config,
            ):
                chunk = ChatCompletionChunk(
                    id=response_id,
                    model=request.model,
                    choices=[
                        ChatCompletionChunkChoice(
                            delta=DeltaMessage(content=text),
                        )
                    ],
                )
                yield f"data: {chunk.model_dump_json()}\n\n"

            # 结束标记
            chunk = ChatCompletionChunk(
                id=response_id,
                model=request.model,
                choices=[
                    ChatCompletionChunkChoice(
                        delta=DeltaMessage(),
                        finish_reason="stop",
                    )
                ],
            )
            yield f"data: {chunk.model_dump_json()}\n\n"
            yield "data: [DONE]\n\n"

        return StreamingResponse(
            generate(),
            media_type="text/event-stream",
        )

    def run(
        self,
        host: str = "0.0.0.0",
        port: int = 8000,
        workers: int = 1,
        **kwargs,
    ) -> None:
        """运行服务"""
        import uvicorn

        uvicorn.run(
            self.app,
            host=host,
            port=port,
            workers=workers,
            **kwargs,
        )


def create_server(
    model: Optional[str] = None,
    device: Optional[str] = None,
    local_dir: Optional[str] = None,
    backend: Optional[BaseLLMBackend] = None,
    dtype: Optional[str] = None,
    attn: Optional[str] = None,
    model_class: Optional[str] = None,
    processor_class: Optional[str] = None,
    vision_processor: Optional[str] = None,
    chat_template_kwargs: Optional[dict] = None,
    lora_modules: Optional[dict] = None,
    compile: bool = False,
    compile_mode: str = "reduce-overhead",
    cache: Optional[str] = None,
    batch_max_size: int = 0,
    batch_wait_s: float = 1.0,
) -> LLMServer:
    """创建 LLM 服务实例

    Args:
        model: 模型 ID
        device: 设备类型
        local_dir: 本地模型目录
        backend: 自定义后端实例
        dtype: 数据类型 (float16/bfloat16/float32)
        attn: 注意力实现 (eager/sdpa/flash_attention_2/flex_attention)
        model_class: 模型类名 (如 "AutoModelForCausalLM")
        processor_class: 处理器类名 (如 "AutoTokenizer")
        vision_processor: 视觉处理器类型 (qwen_vl/general)
        chat_template_kwargs: chat template 额外参数
        lora_modules: LoRA 适配器 {"name": "path", ...}
        compile: 启用 torch.compile
        compile_mode: torch.compile 模式
        cache: KV Cache 实现 (static/dynamic)
        batch_max_size: 最大批次大小，0=禁用批处理
        batch_wait_s: 蓄水池等待时间（秒）

    Returns:
        LLMServer 实例
    """
    return LLMServer(
        backend=backend,
        model=model,
        device=device,
        local_dir=local_dir,
        dtype=dtype,
        attn=attn,
        model_class=model_class,
        processor_class=processor_class,
        vision_processor=vision_processor,
        chat_template_kwargs=chat_template_kwargs,
        lora_modules=lora_modules,
        compile=compile,
        compile_mode=compile_mode,
        cache=cache,
        batch_max_size=batch_max_size,
        batch_wait_s=batch_wait_s,
    )
