import os
import uuid
import requests
from typing import Optional
from contextlib import asynccontextmanager

import torch
import matplotlib
matplotlib.use('Agg') 
import matplotlib.pyplot as plt

from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request, Depends, status
from fastapi.security import APIKeyHeader
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse
from PIL import Image

from sam3.model_builder import build_sam3_image_model
from sam3.model.sam3_image_processor import Sam3Processor
from sam3.visualization_utils import plot_results

# ------------------- 配置与路径 -------------------
STATIC_DIR = "static"
RESULT_IMAGE_DIR = os.path.join(STATIC_DIR, "results")
os.makedirs(RESULT_IMAGE_DIR, exist_ok=True)

# ------------------- API Key 核心配置 (已加固) -------------------
VALID_API_KEY = "123quant-speed"
API_KEY_HEADER_NAME = "X-API-Key"

# 定义 Header 认证
api_key_header = APIKeyHeader(name=API_KEY_HEADER_NAME, auto_error=False)

async def verify_api_key(api_key: Optional[str] = Depends(api_key_header)):
    """
    强制验证 API Key
    """
    # 1. 检查是否有 Key
    if not api_key:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing API Key. Please provide it in the header."
        )
    # 2. 检查 Key 是否正确
    if api_key != VALID_API_KEY:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Invalid API Key."
        )
    # 3. 验证通过
    return True

# ------------------- 生命周期管理 -------------------
@asynccontextmanager
async def lifespan(app: FastAPI):
    print("="*40)
    print("✅ API Key 保护已激活")
    print(f"✅ 有效 Key: {VALID_API_KEY}")
    print("="*40)
    
    print("正在加载 SAM3 模型到 GPU...")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if not torch.cuda.is_available():
        print("警告: 未检测到 GPU，将使用 CPU，速度会较慢。")

    model = build_sam3_image_model()
    model = model.to(device)
    model.eval()

    processor = Sam3Processor(model)

    app.state.model = model
    app.state.processor = processor
    app.state.device = device
    
    print(f"模型加载完成，设备: {device}")
    
    yield 
    
    print("正在清理资源...")

# ------------------- FastAPI 初始化 -------------------
app = FastAPI(
    lifespan=lifespan, 
    title="SAM3 Segmentation API",
    description="## 🔒 受 API Key 保护\n请点击右上角 **Authorize** 并输入: `123quant-speed`",
)

# 手动添加 OpenAPI 安全配置，让 Docs 里的锁头生效
app.openapi_schema = None
def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema
    from fastapi.openapi.utils import get_openapi
    openapi_schema = get_openapi(
        title=app.title,
        version=app.version,
        description=app.description,
        routes=app.routes,
    )
    # 定义安全方案
    openapi_schema["components"]["securitySchemes"] = {
        "APIKeyHeader": {
            "type": "apiKey",
            "in": "header",
            "name": API_KEY_HEADER_NAME,
        }
    }
    # 为所有路径应用安全要求
    for path in openapi_schema["paths"]:
        for method in openapi_schema["paths"][path]:
            openapi_schema["paths"][path][method]["security"] = [{"APIKeyHeader": []}]
    app.openapi_schema = openapi_schema
    return app.openapi_schema

app.openapi = custom_openapi

app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

# ------------------- 辅助函数 -------------------
def load_image_from_url(url: str) -> Image.Image:
    try:
        headers = {'User-Agent': 'Mozilla/5.0'}
        response = requests.get(url, headers=headers, stream=True, timeout=10)
        response.raise_for_status()
        image = Image.open(response.raw).convert("RGB")
        return image
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"无法下载图片: {str(e)}")

def generate_and_save_result(image: Image.Image, inference_state) -> str:
    filename = f"seg_{uuid.uuid4().hex}.jpg"
    save_path = os.path.join(RESULT_IMAGE_DIR, filename)
    plot_results(image, inference_state)
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    return filename

# ------------------- API 接口 (强制依赖验证) -------------------
@app.post("/segment", dependencies=[Depends(verify_api_key)])
async def segment(
    request: Request,
    prompt: str = Form(...),
    file: Optional[UploadFile] = File(None),
    image_url: Optional[str] = Form(None)
):
    if not file and not image_url:
        raise HTTPException(status_code=400, detail="必须提供 file (图片文件) 或 image_url (图片链接)")

    try:
        if file:
            image = Image.open(file.file).convert("RGB")
        elif image_url:
            image = load_image_from_url(image_url)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"图片解析失败: {str(e)}")

    processor = request.app.state.processor

    try:
        inference_state = processor.set_image(image)
        output = processor.set_text_prompt(state=inference_state, prompt=prompt)
        masks, boxes, scores = output["masks"], output["boxes"], output["scores"]
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"模型推理错误: {str(e)}")

    try:
        filename = generate_and_save_result(image, inference_state)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"绘图保存错误: {str(e)}")

    file_url = request.url_for("static", path=f"results/{filename}")

    return JSONResponse(content={
        "status": "success",
        "result_image_url": str(file_url),
        "detected_count": len(masks),
        "scores": scores.tolist() if torch.is_tensor(scores) else scores
    })

if __name__ == "__main__":
    import uvicorn
    # 注意：如果你的文件名不是 fastAPI_nocom.py，请修改下面第一个参数
    uvicorn.run(
        "fastAPI_nocom:app", 
        host="127.0.0.1",
        port=55600,
        proxy_headers=True,
        forwarded_allow_ips="*",
        reload=False # 生产环境建议关闭 reload，确保代码完全重载
    )