# 环境要求：Python 3.8+
# 依赖包：pip install fastapi transformers torch

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict, Any
import torch
from transformers import pipeline
import uvicorn
import time
import logging
from contextlib import asynccontextmanager

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 全局变量存储模型
sentiment_pipeline = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    # 启动时加载模型
    global sentiment_pipeline
    logger.info("开始加载情感分析模型...")
    try:
        # 使用Hugging Face的情感分析模型（更稳定可靠）
        sentiment_pipeline = pipeline(
            "sentiment-analysis",
            model="cardiffnlp/twitter-roberta-base-sentiment-latest",
            device=0 if torch.cuda.is_available() else -1
        )
        logger.info("模型加载完成")
    except Exception as e:
        logger.error(f"模型加载失败: {e}")
        raise e
    
    yield
    
    # 关闭时清理资源
    logger.info("服务关闭")

app = FastAPI(
    title="情感分析服务",
    description="基于BERT的情感分析API",
    version="1.0.0",
    lifespan=lifespan
)

# 请求响应模型
class SentimentRequest(BaseModel):
    text: str
    return_scores: bool = False

class SentimentResponse(BaseModel):
    text: str
    sentiment: str
    confidence: float
    scores: Dict[str, float] = None
    processing_time: float

class BatchSentimentRequest(BaseModel):
    texts: List[str]
    return_scores: bool = False

class BatchSentimentResponse(BaseModel):
    results: List[SentimentResponse]
    total_processing_time: float

# 健康检查端点
@app.get("/")
async def root():
    """服务根路径，返回基本信息"""
    return {
        "message": "情感分析服务正在运行",
        "version": "1.0.0",
        "status": "healthy"
    }

# 健康检查端点
@app.get("/health")
async def health_check():
    """健康检查端点，用于负载均衡器检查"""
    return {"status": "healthy"}

# 单文本情感分析
@app.post("/analyze-sentiment", response_model=SentimentResponse)
async def analyze_sentiment(request: SentimentRequest):
    """
    分析单个文本的情感倾向
    
    Args:
        request: 包含文本和选项的请求
        
    Returns:
        情感分析结果
    """
    if sentiment_pipeline is None:
        raise HTTPException(status_code=503, detail="模型未加载")
    
    start_time = time.time()
    
    try:
        # 执行情感分析
        result = sentiment_pipeline(request.text)
        
        processing_time = time.time() - start_time
        
        # 解析结果
        sentiment_label = result[0]['label']
        confidence = result[0]['score']
        
        # 将英文标签转换为中文
        label_mapping = {
            'LABEL_0': '负面',
            'LABEL_1': '中性', 
            'LABEL_2': '正面'
        }
        chinese_label = label_mapping.get(sentiment_label, sentiment_label)
        
        response_data = {
            "text": request.text,
            "sentiment": chinese_label,
            "confidence": confidence,
            "processing_time": processing_time
        }
        
        # 如果需要返回详细分数
        if request.return_scores:
            scores = {}
            for item in result:
                label = label_mapping.get(item['label'], item['label'])
                scores[label] = item['score']
            response_data["scores"] = scores
        
        return SentimentResponse(**response_data)
        
    except Exception as e:
        logger.error(f"情感分析失败: {e}")
        raise HTTPException(status_code=500, detail=f"分析失败: {str(e)}")

# 批量情感分析
@app.post("/analyze-sentiment-batch", response_model=BatchSentimentResponse)
async def analyze_sentiment_batch(request: BatchSentimentRequest):
    """
    批量分析多个文本的情感倾向
    
    Args:
        request: 包含文本列表的请求
        
    Returns:
        批量情感分析结果
    """
    if sentiment_pipeline is None:
        raise HTTPException(status_code=503, detail="模型未加载")
    
    start_time = time.time()
    results = []
    
    try:
        # 批量处理文本
        batch_results = sentiment_pipeline(request.texts)
        
        # 标签映射
        label_mapping = {
            'LABEL_0': '负面',
            'LABEL_1': '中性', 
            'LABEL_2': '正面'
        }
        
        for i, (text, result) in enumerate(zip(request.texts, batch_results)):
            sentiment_label = result['label']
            confidence = result['score']
            chinese_label = label_mapping.get(sentiment_label, sentiment_label)
            
            response_data = {
                "text": text,
                "sentiment": chinese_label,
                "confidence": confidence,
                "processing_time": 0  # 批量处理中单个时间难以计算
            }
            
            if request.return_scores:
                # 对于批量处理，这里简化处理
                response_data["scores"] = {label_mapping.get(result['label'], result['label']): result['score']}
            
            results.append(SentimentResponse(**response_data))
        
        total_processing_time = time.time() - start_time
        
        return BatchSentimentResponse(
            results=results,
            total_processing_time=total_processing_time
        )
        
    except Exception as e:
        logger.error(f"批量情感分析失败: {e}")
        raise HTTPException(status_code=500, detail=f"批量分析失败: {str(e)}")

# 模型信息端点
@app.get("/model-info")
async def get_model_info():
    """获取当前加载的模型信息"""
    if sentiment_pipeline is None:
        raise HTTPException(status_code=503, detail="模型未加载")
    
    return {
        "model_name": sentiment_pipeline.model.config.name_or_path,
        "device": str(sentiment_pipeline.device),
        "status": "loaded",
        "model_type": "RoBERTa",
        "source": "Hugging Face"
    }

if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=False  # 生产环境建议关闭热重载
    )