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

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import AsyncGenerator
from typing import List, Dict, Any
import torch
import asyncio
import json
from prometheus_client import Counter, Histogram, generate_latest, CollectorRegistry, REGISTRY
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
)

# Prometheus监控指标
# 创建自定义注册表避免冲突
custom_registry = CollectorRegistry()

# 定义指标
REQUEST_COUNT = Counter(
    'sentiment_api_requests_total', 
    'Total sentiment analysis requests', 
    ['method', 'endpoint'],
    registry=custom_registry
)

REQUEST_DURATION = Histogram(
    'sentiment_api_request_duration_seconds', 
    'Sentiment analysis request duration', 
    ['method', 'endpoint'],
    registry=custom_registry
)

# 请求响应模型
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

# 流式文本生成函数
async def generate_text_stream(prompt: str, max_length: int, temperature: float) -> AsyncGenerator[str, None]:
    """流式文本生成"""
    if sentiment_pipeline is None:
        yield "模型未加载"
        return
    
    try:
        # 使用pipeline生成文本
        result = sentiment_pipeline(
            prompt,
            max_length=max_length,
            temperature=temperature,
            do_sample=True,
            pad_token_id=sentiment_pipeline.tokenizer.eos_token_id
        )
        
        # 提取生成的文本（去除原始prompt）
        generated_text = result[0]['generated_text']
        new_text = generated_text[len(prompt):].strip()
        
        # 模拟流式输出（将文本分块发送）
        words = new_text.split()
        for i, word in enumerate(words):
            await asyncio.sleep(0.1)  # 模拟生成延迟
            yield word + " "
            
    except Exception as e:
        logger.error(f"文本生成失败: {e}")
        yield f"生成失败: {str(e)}"

# 健康检查端点
@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.post("/generate-text-stream")
async def generate_text_stream_endpoint(request: SentimentRequest):
    """
    流式文本生成端点
    
    Args:
        request: 文本生成请求参数
        
    Returns:
        流式响应，每个文本片段作为SSE事件
    """
    if sentiment_pipeline is None:
        raise HTTPException(status_code=503, detail="模型未加载")
    
    start_time = time.time()
    
    async def generate():
        try:
            async for chunk in generate_text_stream(
                request.prompt, 
                request.max_length, 
                request.temperature
            ):
                # 发送SSE格式的数据
                yield f"data: {json.dumps({'text': chunk, 'timestamp': time.time()})}\n\n"
                
        except Exception as e:
            logger.error(f"流式生成失败: {e}")
            yield f"data: {json.dumps({'error': str(e)})}\n\n"
    
    return SentimentResponse(
        generate(),
        media_type="text/plain",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        }
    )

# 中间件
@app.middleware("http")
async def monitor_requests(request, call_next):
    start_time = time.time()
    REQUEST_COUNT.labels(method=request.method, endpoint=request.url.path).inc()
    
    response = await call_next(request)
    
    duration = time.time() - start_time
    REQUEST_DURATION.labels(method=request.method, endpoint=request.url.path).observe(duration)
    
    return response

# 监控指标端点
@app.get("/metrics")
async def metrics():
    """Prometheus监控指标端点"""
    return generate_latest(custom_registry)

# 模型信息端点
@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  # 生产环境建议关闭热重载
    )
