#!/usr/bin/env python3
"""WeMM-Embedding 独立服务 - FastAPI"""
import asyncio
import logging
import os
import tempfile
import base64
from typing import List, Optional, Union
from pathlib import Path

import torch
import numpy as np
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# 尝试导入 WeMM-Embedding
try:
    from sentence_transformers import SentenceTransformer
    _WEMM_AVAILABLE = True
except ImportError:
    _WEMM_AVAILABLE = False

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ========== 配置 ==========
MODEL_ID = os.getenv("WEMM_MODEL", "tencent/WeMM-Embedding-2B")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MAX_DIMENSION = int(os.getenv("MAX_DIMENSION", "2048"))
PORT = int(os.getenv("PORT", "8765"))


# ========== 数据模型 ==========
class EmbeddingRequest(BaseModel):
    """嵌入请求"""
    inputs: Union[str, List[Union[str, dict]]]
    dimension: Optional[int] = Field(None, ge=64, le=MAX_DIMENSION, description="输出维度 (64-4096)")
    normalize: bool = True


class EmbeddingResponse(BaseModel):
    """嵌入响应"""
    embeddings: List[List[float]]
    model: str
    dimensions: int


class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    device: str
    cuda_available: bool


# ========== 服务类 ==========
class WeMMService:
    """WeMM-Embedding 服务封装"""
    
    def __init__(self, model_id: str, device: str):
        self.model_id = model_id
        self.device = device
        self.model = None
        
    async def load(self):
        """异步加载模型"""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._load_sync)
    
    def _load_sync(self):
        if _WEMM_AVAILABLE:
            self.model = SentenceTransformer(
                self.model_id,
                device=self.device,
                trust_remote_code=True
            )
            logger.info(f"✅ 模型加载完成: {self.model_id} ({self.device})")
            return True
        return False
    
    async def encode(
        self, 
        inputs: Union[str, List[Union[str, dict]]],
        dimension: Optional[int] = None,
        normalize: bool = True
    ) -> np.ndarray:
        """编码输入"""
        if self.model is None:
            raise RuntimeError("模型未加载")
        
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None, 
            lambda: self._encode_sync(inputs, dimension, normalize)
        )
    
    def _encode_sync(
        self, 
        inputs: Union[str, List[Union[str, dict]]],
        dimension: Optional[int],
        normalize: bool
    ) -> np.ndarray:
        args = {}
        if dimension:
            args['truncate_dim'] = dimension
        if not normalize:
            args['normalize_embeddings'] = False
            
        return self.model.encode(inputs, **args)
    
    def is_ready(self) -> bool:
        return self.model is not None


# ========== FastAPI App ==========
app = FastAPI(
    title="WeMM-Embedding Service",
    description="腾讯 WeMM-Embedding 多模态嵌入服务",
    version="1.0.0"
)

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

# 全局服务实例
service: Optional[WeMMService] = None


@app.on_event("startup")
async def startup_event():
    global service
    logger.info(f"🚀 启动 WeMM-Embedding 服务...")
    service = WeMMService(MODEL_ID, DEVICE)
    loaded = await service.load()
    if not loaded:
        logger.error("❌ 模型加载失败，请检查依赖")


@app.get("/health", response_model=HealthResponse)
async def health():
    return HealthResponse(
        status="ok" if service and service.is_ready() else "error",
        model_loaded=service.is_ready() if service else False,
        device=DEVICE,
        cuda_available=torch.cuda.is_available()
    )


@app.post("/embed", response_model=EmbeddingResponse)
async def embed(request: EmbeddingRequest):
    if not service or not service.is_ready():
        raise HTTPException(status_code=503, detail="模型未加载")
    
    try:
        embeddings = await service.encode(
            request.inputs,
            dimension=request.dimension,
            normalize=request.normalize
        )
        
        return EmbeddingResponse(
            embeddings=embeddings.tolist(),
            model=MODEL_ID,
            dimensions=embeddings.shape[-1] if len(embeddings.shape) > 1 else 1
        )
    except Exception as e:
        logger.error(f"编码失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/embed/image")
async def embed_image(
    file: UploadFile = File(...),
    query: str = "",
    dimension: int = MAX_DIMENSION
):
    """上传图像文件并获取嵌入"""
    if not service or not service.is_ready():
        raise HTTPException(status_code=503, detail="模型未加载")
    
    try:
        contents = await file.read()
        
        # 保存到临时文件
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
            tmp.write(contents)
            tmp_path = tmp.name
        
        try:
            inputs = [{"type": "image", "image": tmp_path}]
            if query:
                inputs.append({"type": "text", "text": query})
            
            embeddings = await service.encode(inputs, dimension=dimension)
            
            return {
                "embedding": embeddings[0].tolist(),
                "dimensions": embeddings.shape[-1],
                "file": file.filename
            }
        finally:
            os.unlink(tmp_path)
    except Exception as e:
        logger.error(f"图像编码失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/embed/video")
async def embed_video(
    file: UploadFile = File(...),
    query: str = "",
    dimension: int = MAX_DIMENSION
):
    """上传视频文件并获取嵌入（采样64帧）"""
    if not service or not service.is_ready():
        raise HTTPException(status_code=503, detail="模型未加载")
    
    try:
        contents = await file.read()
        
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
            tmp.write(contents)
            tmp_path = tmp.name
        
        try:
            inputs = [{"type": "video", "video": tmp_path}]
            if query:
                inputs.append({"type": "text", "text": query})
            
            embeddings = await service.encode(inputs, dimension=dimension)
            
            return {
                "embedding": embeddings[0].tolist(),
                "dimensions": embeddings.shape[-1],
                "file": file.filename
            }
        finally:
            os.unlink(tmp_path)
    except Exception as e:
        logger.error(f"视频编码失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/similarity")
async def similarity(request: dict):
    """计算相似度"""
    if not service or not service.is_ready():
        raise HTTPException(status_code=503, detail="模型未加载")
    
    try:
        texts = request.get("texts", [])
        dimension = request.get("dimension", MAX_DIMENSION)
        
        if not texts:
            raise HTTPException(status_code=400, detail="需要提供 texts 列表")
        
        embeddings = await service.encode(texts, dimension=dimension)
        
        # 计算余弦相似度
        norm = np.linalg.norm(embeddings, axis=-1, keepdims=True)
        normalized = embeddings / norm
        sim_matrix = np.dot(normalized, normalized.T)
        
        return {
            "similarities": sim_matrix.tolist(),
            "labels": texts
        }
    except Exception as e:
        logger.error(f"相似度计算失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn
    logger.info(f"🎯 启动服务: http://0.0.0.0:{PORT}")
    logger.info(f"📊 模型: {MODEL_ID}")
    logger.info(f"💻 设备: {DEVICE}")
    
    uvicorn.run(
        "wemm_service:app",
        host="0.0.0.0",
        port=PORT,
        log_level="info"
    )