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 numpy as np
from threading import Lock
import uvicorn

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.device = torch.device("cpu")
        self.model = None
        self.tokenizer = None
        self.lock = Lock()
        self.setup_logging()
        torch.set_num_threads(4)  # CPU优化

    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
            ).to(self.device)
            self.model.eval()
            torch.set_grad_enabled(False)
            self.logger.info(f"模型加载成功，使用设备: {self.device}")
        except Exception as e:
            self.logger.error(f"模型初始化失败: {str(e)}")
            raise

    @lru_cache(maxsize=1000)
    def get_embedding(self, text: str) -> List[float]:
        """同步生成嵌入向量，带缓存"""
        with self.lock:
            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

embedding_service = EmbeddingService()
app = FastAPI()

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

@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:
        # 使用run_in_executor避免事件循环问题
        embedding = await asyncio.get_running_loop().run_in_executor(
            None,
            embedding_service.get_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)
    }

def gradio_interface(text: str) -> Dict:
    try:
        embedding = embedding_service.get_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())
    gr.mount_gradio_app(app, iface, path="/ui")
    uvicorn.run(app, host="0.0.0.0", port=7860, workers=1)