import uuid
import time
import redis
import requests
import json
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
# 初始化Redis（用于存储幂等状态、缓存结果、分布式锁）
redis_client = redis.Redis(
    host="localhost",
    port=6379,
    db=0,
    decode_responses=True
)
app = FastAPI(title="大模型接口幂等性设计Demo")
# 本地模型API地址
LOCAL_MODEL_API = "http://192.168.3.6:8000/v1/chat/completions"
# 模拟大模型请求体
class ChatRequest(BaseModel):
    session_id: str  # 会话ID
    prompt: str      # 用户提问内容
# ===================== 核心工具函数 =====================
def get_distributed_lock(key: str, expire: int = 30):
    """获取分布式锁，防止并发重复执行"""
    return redis_client.set(key + "_lock", "locked", nx=True, ex=expire)
def release_distributed_lock(key: str):
    """释放锁"""
    redis_client.delete(key + "_lock")
def call_local_model(prompt: str):
    """调用本地大模型API"""
    # 尝试多种请求格式
    payloads = [
        # 格式1: OpenAI标准格式
        {
            "model": "default",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.7,
            "max_tokens": 2048
        },
        # 格式2: 简化格式（部分本地模型）
        {
            "prompt": prompt,
            "max_tokens": 2048,
            "temperature": 0.7
        },
        # 格式3: 带system的完整格式
        {
            "model": "default",
            "messages": [
                {"role": "system", "content": "你是一个 helpful assistant."},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.7,
            "max_tokens": 2048
        }
    ]
    
    for idx, payload in enumerate(payloads):
        start_time = time.time()
        try:
            # print(f"[INFO] 尝试请求格式 {idx + 1}: {list(payload.keys())}")
            response = requests.post(
                LOCAL_MODEL_API,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=120
            )
            latency = round(time.time() - start_time, 3)
            
            # print(f"[<-] 响应状态码: {response.status_code}")
            
            if response.status_code == 200:
                result = response.json()
                # print(f"[INFO] 请求成功，格式 {idx + 1} 兼容")
                
                # 提取生成的文本
                generated_text = ""
                if "choices" in result and len(result["choices"]) > 0:
                    choice = result["choices"][0]
                    if "text" in choice:
                        generated_text = choice["text"]
                    elif "delta" in choice and "content" in choice["delta"]:
                        generated_text = choice["delta"]["content"]
                    elif "message" in choice and "content" in choice["message"]:
                        generated_text = choice["message"]["content"]
                
                return {
                    "success": True,
                    "content": generated_text,
                    "latency": latency,
                    "raw_response": result
                }
            else:
                # print(f"[WARN] 格式 {idx + 1} 失败: {response.status_code}")
                if idx == len(payloads) - 1:
                    return {
                        "success": False,
                        "error": f"模型请求失败: {response.status_code}, 响应: {response.text[:200]}",
                        "latency": latency
                    }
        except Exception as e:
            # print(f"[WARN] 格式 {idx + 1} 异常: {str(e)}")
            if idx == len(payloads) - 1:
                return {
                    "success": False,
                    "error": f"模型请求异常: {str(e)}",
                    "latency": 0
                }
def simulate_billing():
    """模拟计费逻辑（仅执行一次）"""
    print("【计费】执行一次扣费")
    return True
# ===================== 幂等性核心接口 =====================
@app.post("/v1/chat", summary="大模型对话接口（幂等版）")
async def chat_completions(
    request: ChatRequest,
    idempotency_key: str = Header(..., alias="Idempotency-Key")
):
    """
    幂等核心：
    1. 从请求头获取唯一幂等键 Idempotency-Key
    2. 重复请求直接返回缓存结果，不推理、不扣费
    3. 支持会话防重 + 请求防重双重保障
    """
    session_id = request.session_id
    prompt = request.prompt
    # ========== 第一步：校验幂等ID是否存在 ==========
    if not idempotency_key:
        raise HTTPException(status_code=400, detail="缺少幂等标识 Idempotency-Key")
    # 组合唯一键：幂等ID + 会话ID（双重防重）
    unique_key = f"idem:{idempotency_key}:{session_id}"
    # ========== 第二步：查询请求状态 ==========
    cached_data = redis_client.hgetall(unique_key)
    # 情况1：已完成，直接返回缓存结果
    if cached_data and cached_data.get("status") == "SUCCESS":
        print(f"【幂等】重复请求，直接返回缓存结果 | Key: {idempotency_key}")
        return {
            "code": 200,
            "message": "重复请求，返回历史结果",
            "data": cached_data.get("answer"),
            "is_retry": True
        }
    # 情况2：处理中，拒绝重复执行
    if cached_data and cached_data.get("status") == "PROCESSING":
        raise HTTPException(status_code=429, detail="请求正在推理中，请稍后重试")
    # ========== 第三步：获取分布式锁，防止并发 ==========
    lock_acquired = get_distributed_lock(unique_key)
    if not lock_acquired:
        raise HTTPException(status_code=429, detail="并发请求被拒绝，保证幂等执行")
    try:
        # 标记状态：处理中
        redis_client.hset(unique_key, "status", "PROCESSING")
        redis_client.hset(unique_key, "prompt", prompt)
        redis_client.hset(unique_key, "session_id", session_id)
        # 过期时间：1小时，避免数据堆积
        redis_client.expire(unique_key, 3600)
        # ========== 第四步：执行核心逻辑 ==========
        # 1. 大模型推理（仅执行一次）
        model_result = call_local_model(prompt)
        if not model_result["success"]:
            raise HTTPException(status_code=500, detail=model_result["error"])
        answer = model_result["content"]
        # print(f"[INFO] 模型生成内容长度: {len(answer)} 字符")
        
        # 2. 计费（仅执行一次）
        simulate_billing()
        # ========== 第五步：更新状态为成功，缓存结果 ==========
        redis_client.hset(unique_key, "status", "SUCCESS")
        redis_client.hset(unique_key, "answer", answer)
        redis_client.hset(unique_key, "prompt", prompt)
        return {
            "code": 200,
            "message": "请求成功",
            "data": answer,
            "is_retry": False
        }
    except Exception as e:
        # 失败：更新状态
        import traceback
        error_msg = f"推理失败: {str(e)}"
        # print(f"[ERROR] {error_msg}")
        # print(f"[ERROR] 详细堆栈:\n{traceback.format_exc()}")
        redis_client.hset(unique_key, "status", "FAILED")
        raise HTTPException(status_code=500, detail=error_msg)
    finally:
        # 释放分布式锁
        release_distributed_lock(unique_key)
# 测试接口：生成幂等Key（给客户端使用）
@app.get("/v1/idempotency/generate", summary="生成幂等唯一Key")
def generate_idempotency_key():
    return {"idempotency_key": str(uuid.uuid4())}
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

# ===================== 测试脚本 =====================
def test_idempotency():
    """幂等性测试"""
    import requests
    import time
    
    BASE_URL = "http://localhost:8000"
    
    # 1. 生成幂等Key
    print("=" * 50)
    print("【测试1】生成幂等Key")
    resp = requests.get(f"{BASE_URL}/v1/idempotency/generate")
    key = resp.json()["idempotency_key"]
    print(f"生成Key: {key}")
    
    # 2. 首次请求
    print("\n" + "=" * 50)
    print("【测试2】首次请求（应执行推理+计费）")
    headers = {
        "Content-Type": "application/json",
        "Idempotency-Key": key
    }
    data = {
        "session_id": "session_001",
        "prompt": "你好,随意回复20个字"
    }
    start = time.time()
    resp = requests.post(f"{BASE_URL}/v1/chat", headers=headers, json=data)
    print(f"耗时: {time.time()-start:.2f}s")
    print(f"响应: {resp.json()}")
    
    # 3. 重复请求（幂等）
    print("\n" + "=" * 50)
    print("【测试3】重复请求（应直接返回缓存，无计费）")
    start = time.time()
    resp = requests.post(f"{BASE_URL}/v1/chat", headers=headers, json=data)
    print(f"耗时: {time.time()-start:.2f}s")
    print(f"响应: {resp.json()}")
    
    # 4. 不同Key请求
    print("\n" + "=" * 50)
    print("【测试4】不同Key请求（应独立执行）")
    resp = requests.get(f"{BASE_URL}/v1/idempotency/generate")
    new_key = resp.json()["idempotency_key"]
    headers["Idempotency-Key"] = new_key
    start = time.time()
    resp = requests.post(f"{BASE_URL}/v1/chat", headers=headers, json=data)
    print(f"耗时: {time.time()-start:.2f}s")
    print(f"响应: {resp.json()}")
    
    print("\n" + "=" * 50)
    print("测试完成！")
test_idempotency()