#!/usr/bin/env python3
"""
vLLM API 服务器（macOS 替代方案）

使用 vLLM 的核心引擎和 Uvicorn ASGI 框架，
在 macOS 上实现 OpenAI 兼容 API 服务。

原生 vLLM serve 命令在 macOS 上不可用（C 扩展兼容性问题），
本方案使用 vLLM 的 API 类和手动 HTTP 路由实现等效功能。
"""

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
import uvicorn
from vllm import LLM, SamplingParams
from vllm.entrypoints.openai.protocol import (
    ChatCompletionRequest,
    CompletionRequest,
)
import json
import asyncio
from typing import AsyncGenerator

# ============================================================================
# vLLM 引擎初始化（全局变量，在 main 中初始化）
# ============================================================================

llm = None

# ============================================================================
# FastAPI 应用初始化
# ============================================================================

app = FastAPI(title="vLLM OpenAI-Compatible API")

# ============================================================================
# 健康检查端点
# ============================================================================

@app.get("/health")
async def health():
    """健康检查端点"""
    return {"status": "ok", "model": "Qwen/Qwen2.5-1.5B-Instruct"}

# ============================================================================
# 聊天补全端点（OpenAI 兼容）
# ============================================================================

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    """
    聊天补全 API（OpenAI 兼容）
    
    请求格式：
    {
      "messages": [
        {"role": "user", "content": "你好"}
      ],
      "temperature": 0.7,
      "max_tokens": 256,
      "stream": false
    }
    """
    
    # 从消息列表构建 prompt
    # vLLM 提供了 apply_chat_template 方法
    prompt = llm.get_tokenizer().apply_chat_template(
        request.messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    
    # 配置生成参数
    # vLLM 的 SamplingParams 管理温度、top-p 等采样参数
    sampling_params = SamplingParams(
        temperature=request.temperature,
        top_p=request.top_p if request.top_p else 0.95,
        max_tokens=request.max_tokens if request.max_tokens else 256,
    )
    
    # 如果请求流式输出
    if request.stream:
        # 返回流式响应
        async def generate():
            # vLLM 的 generate 方法支持流式输出
            outputs = llm.generate([prompt], sampling_params)
            output = outputs[0]
            text = output.outputs[0].text
            
            # 流式返回 OpenAI 格式的消息
            for char in text:
                chunk = {
                    "choices": [{
                        "delta": {"content": char},
                        "index": 0,
                    }]
                }
                yield f"data: {json.dumps(chunk)}\n\n"
            
            # 流式结束标记
            yield "data: [DONE]\n\n"
        
        return StreamingResponse(generate(), media_type="text/event-stream")
    
    else:
        # 非流式输出：直接生成，一次返回完整结果
        outputs = llm.generate([prompt], sampling_params)
        output = outputs[0]
        text = output.outputs[0].text
        
        return JSONResponse({
            "choices": [{
                "message": {
                    "role": "assistant",
                    "content": text,
                },
                "index": 0,
                "finish_reason": "stop",
            }]
        })

# ============================================================================
# 文本补全端点（OpenAI 兼容）
# ============================================================================

@app.post("/v1/completions")
async def completions(request: CompletionRequest):
    """
    文本补全 API（OpenAI 兼容）
    
    请求格式：
    {
      "prompt": "今天天气",
      "temperature": 0.7,
      "max_tokens": 128
    }
    """
    
    # 配置生成参数
    sampling_params = SamplingParams(
        temperature=request.temperature,
        top_p=request.top_p if request.top_p else 0.95,
        max_tokens=request.max_tokens if request.max_tokens else 256,
    )
    
    # 生成文本
    outputs = llm.generate([request.prompt], sampling_params)
    output = outputs[0]
    text = output.outputs[0].text
    
    return JSONResponse({
        "choices": [{
            "text": text,
            "index": 0,
            "finish_reason": "stop",
        }]
    })

# ============================================================================
# 启动服务器
# ============================================================================

if __name__ == "__main__":
    """
    启动 Uvicorn ASGI 服务器
    
    Uvicorn 是一个轻量级 ASGI 服务器，比 Flask 更高效，
    原生支持异步处理和流式响应。
    
    启动命令：
      python main.py
    
    生产部署（多 worker）：
      uvicorn main:app --host 0.0.0.0 --port 8000 --workers 1
    
    注意：
    - workers=1 是因为 vLLM 引擎不是线程安全的
    - 如果需要并发，应该使用进程池或消息队列
    - macOS 上必须在 if __name__ == '__main__': 块中初始化 vLLM，因为它使用多进程
    """
    # vLLM 使用 LLM 引擎作为核心推理组件
    # 相比 Transformers，vLLM 提供了自动 Batch 处理和 KV Cache 管理
    print("Initializing vLLM engine...")
    llm = LLM(
        model="Qwen/Qwen2.5-1.5B-Instruct",
        dtype="bfloat16",  # macOS 推荐 bfloat16
        gpu_memory_utilization=0.3,  # 限制内存使用（针对 CPU 推理）
        max_model_len=2048,  # 限制最大序列长度，节省内存
    )
    print("vLLM engine initialized!\n")
    
    print("Starting vLLM API server on http://0.0.0.0:8000")
    print("API 文档：http://localhost:8000/docs")
    
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        log_level="info",
    )