import asyncio
import logging
import torch
import gradio as gr
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict
from functools import lru_cache
import uvicorn
import psutil
import numpy as np

class EmbeddingRequest(BaseModel):
    input: str
    model: str = "jinaai/jina-embeddings-v3"

class EmbeddingResponse(BaseModel):
    status: str
    embeddings: List[List[float]]

class EmbeddingService:
    def __init__(self):
        self.model_name = "jinaai/jina-embeddings-v3"
        self.max_length = 512
        self.batch_size = 8
        self.device = torch.device("cpu")
        self.num_threads = min(psutil.cpu_count(), 4)  # 限制CPU线程数
        self.model = None
        self.tokenizer = None
        self.setup_logging()
        
        # CPU优化配置
        torch.set_num_threads(self.num_threads)

    def setup_logging(self):
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s'
        )
        self.logger = logging.getLogger(__name__)

    async def initialize(self):
        try:
            from transformers import AutoTokenizer, AutoModel
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name,
                trust_remote_code=True
            )
            self.model = AutoModel.from_pretrained(
                self.model_name,
                trust_remote_code=True,
                torch_dtype=torch.float32  # CPU使用float32
            ).to(self.device)
            
            self.model.eval()
            torch.set_grad_enabled(False)
            self.logger.info(f"模型加载成功，CPU线程数: {self.num_threads}")
        except Exception as e:
            self.logger.error(f"模型初始化失败: {str(e)}")
            raise

    @lru_cache(maxsize=1000)
    async def generate_embedding(self, text: str) -> List[float]:
        try:
            inputs = self.tokenizer(
                text,
                return_tensors="pt",
                truncation=True,
                max_length=self.max_length,
                padding=True
            )

            with torch.no_grad():
                outputs = self.model(**inputs).last_hidden_state.mean(dim=1)
                return outputs.numpy().tolist()[0]
        except Exception as e:
            self.logger.error(f"生成嵌入向量失败: {str(e)}")
            raise

# FastAPI应用初始化
app = FastAPI(
    title="Jina Embeddings API",
    description="Text embedding generation service using jina-embeddings-v3",
    version="1.0.0"
)

# 初始化服务
embedding_service = EmbeddingService()

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

# API端点
@app.post("/generate_embeddings", response_model=EmbeddingResponse)
@app.post("/api/v1/embeddings", response_model=EmbeddingResponse)
@app.post("/hf/v1/embeddings", response_model=EmbeddingResponse)
@app.post("/api/v1/chat/completions", response_model=EmbeddingResponse)
@app.post("/hf/v1/chat/completions", response_model=EmbeddingResponse)
async def generate_embeddings(request: EmbeddingRequest):
    try:
        embedding = await embedding_service.generate_embedding(request.input)
        return EmbeddingResponse(
            status="success",
            embeddings=[embedding]
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {
        "status": "active",
        "model": embedding_service.model_name,
        "device": str(embedding_service.device),
        "cpu_threads": embedding_service.num_threads
    }

# Gradio界面
def gradio_interface(text: str) -> Dict:
    try:
        embedding = asyncio.run(embedding_service.generate_embedding(text))
        return {
            "status": "success",
            "embeddings": [embedding]
        }
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(lines=3, label="输入文本"),
    outputs=gr.JSON(label="嵌入向量结果"),
    title="Jina Embeddings V3",
    description="使用jina-embeddings-v3模型生成文本嵌入向量",
    examples=[["这是一个测试句子。"]]
)

@app.on_event("startup")
async def startup_event():
    await embedding_service.initialize()

if __name__ == "__main__":
    # 初始化服务
    asyncio.run(embedding_service.initialize())
    
    # 挂载Gradio应用
    gr.mount_gradio_app(app, iface, path="/ui")
    
    # 启动服务
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860,
        workers=1,
        loop="asyncio"
    )