#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
VLM 图生文推理服务 — AndesVL-1B (千问体系)
为 RDK X5 提供 OpenAI 兼容 API（纯 CPU 推理）

启动方式: python vlm_server.py
API 地址: http://192.168.10.210:8000
"""

import base64
import io
import time
import uuid
import logging
import sys
from contextlib import asynccontextmanager

import torch
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from PIL import Image
from transformers import AutoModel, AutoTokenizer, AutoImageProcessor

# ========== 配置 ==========
MODEL_PATH = "/home/root/models/AndesVL-1B-Instruct"   # RDK X5 上模型路径
HOST = "0.0.0.0"
PORT = 8000

# ========== 日志 ==========
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("vlm_server")

# ========== 全局变量 ==========
model = None
tokenizer = None
image_processor = None


def load_model():
    """加载 AndesVL-1B 模型"""
    global model, tokenizer, image_processor

    logger.info(f"模型路径: {MODEL_PATH}")
    logger.info("加载 tokenizer & image processor...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
    image_processor = AutoImageProcessor.from_pretrained(
        MODEL_PATH, trust_remote_code=True
    )

    logger.info("加载 AndesVL-1B 模型 (fp32 CPU, 约 18 秒)...")
    t0 = time.time()
    model = AutoModel.from_pretrained(
        MODEL_PATH,
        trust_remote_code=True,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True,
    )
    model.eval()
    logger.info(f"模型加载完成，耗时 {time.time() - t0:.1f}s")
    logger.info("AndesVL-1B 推理服务就绪！")


@asynccontextmanager
async def lifespan(app: FastAPI):
    load_model()
    yield


# ========== FastAPI 应用 ==========
app = FastAPI(title="VLM Server (AndesVL-1B)", lifespan=lifespan)

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


# ========== Pydantic 模型 (OpenAI 格式) ==========
class ImageUrl(BaseModel):
    url: str


class ContentPart(BaseModel):
    type: str
    text: str | None = None
    image_url: ImageUrl | None = None


class Message(BaseModel):
    role: str
    content: str | list[ContentPart]


class ChatCompletionRequest(BaseModel):
    model: str
    messages: list[Message]
    max_tokens: int = 100
    temperature: float | None = None


# ========== 辅助函数 ==========
def decode_base64_image(data_url: str) -> Image.Image:
    """从 data URL 解码图像"""
    if "," in data_url:
        base64_str = data_url.split(",", 1)[1]
    else:
        base64_str = data_url
    image_bytes = base64.b64decode(base64_str)
    return Image.open(io.BytesIO(image_bytes)).convert("RGB")


def build_messages(messages: list[Message]) -> list[dict]:
    """OpenAI 格式 → AndesVL chat 格式"""
    result = []
    for msg in messages:
        if isinstance(msg.content, str):
            result.append({"role": msg.role, "content": msg.content})
        else:
            content_list = []
            for part in msg.content:
                if part.type == "text" and part.text:
                    content_list.append({"type": "text", "text": part.text})
                elif part.type == "image_url" and part.image_url:
                    image = decode_base64_image(part.image_url.url)
                    content_list.append({"type": "image", "image": image})
            result.append({"role": msg.role, "content": content_list})
    return result


# ========== API 路由 ==========
@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [
            {
                "id": "./OpenGVLab/InternVL3-1B/",
                "object": "model",
                "created": 1700000000,
                "owned_by": "local",
            }
        ],
    }


@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    try:
        messages = build_messages(request.messages)
        logger.info(f"收到请求, model={request.model}, max_tokens={request.max_tokens}")

        start_time = time.time()
        response_text = model.chat(
            messages,
            tokenizer,
            image_processor,
            max_new_tokens=request.max_tokens,
        )
        elapsed = time.time() - start_time
        logger.info(f"推理完成，耗时 {elapsed:.1f}s, 结果: {response_text[:50]}...")

        return {
            "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": request.model,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": response_text},
                    "finish_reason": "stop",
                }
            ],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
            },
        }

    except Exception as e:
        logger.error(f"推理出错: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health():
    return {"status": "ok", "model_loaded": model is not None}


# ========== 入口 ==========
if __name__ == "__main__":
    logger.info("=" * 50)
    logger.info("VLM 图生文推理服务启动中...")
    logger.info(f"监听: http://{HOST}:{PORT}")
    logger.info("=" * 50)
    uvicorn.run(app, host=HOST, port=PORT, log_level="info")