import os
import uuid
import time
import requests
import numpy as np
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 crop_and_save_objects(image: Image.Image, masks, boxes, output_dir: str = RESULT_IMAGE_DIR) -> list[str]:
    """
    根据 mask 和 box 裁剪出独立的对象图片 (保留透明背景)
    """
    saved_files = []
    # Convert image to numpy array
    img_arr = np.array(image) # RGB (H, W, 3)
    
    for i, (mask, box) in enumerate(zip(masks, boxes)):
        # Handle tensor/numpy conversions
        if isinstance(mask, torch.Tensor):
            mask_np = mask.cpu().numpy().squeeze()
        else:
            mask_np = mask.squeeze()
            
        if isinstance(box, torch.Tensor):
            box_np = box.cpu().numpy()
        else:
            box_np = box
            
        # Get coordinates
        x1, y1, x2, y2 = map(int, box_np)
        
        # Ensure coordinates are within bounds
        x1 = max(0, x1)
        y1 = max(0, y1)
        x2 = min(image.width, x2)
        y2 = min(image.height, y2)
        
        # Check valid crop
        if x2 <= x1 or y2 <= y1:
            continue
            
        # Create Alpha channel from mask (0 or 255)
        # mask_np is boolean or float 0..1. If boolean, *255 -> 0/255.
        alpha = (mask_np * 255).astype(np.uint8)
        
        # Combine RGB and Alpha
        rgba = np.dstack((img_arr, alpha))
        
        # Convert back to PIL for cropping
        pil_rgba = Image.fromarray(rgba)
        
        # Crop to bounding box
        cropped = pil_rgba.crop((x1, y1, x2, y2))
        
        # Save
        filename = f"tarot_{uuid.uuid4().hex}_{i}.png" # Use png for transparency
        save_path = os.path.join(output_dir, filename)
        cropped.save(save_path)
        saved_files.append(filename)
        
    return saved_files

def generate_and_save_result(image: Image.Image, inference_state, output_dir: str = RESULT_IMAGE_DIR) -> str:
    filename = f"seg_{uuid.uuid4().hex}.jpg"
    save_path = os.path.join(output_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
    })

@app.post("/segment_tarot", dependencies=[Depends(verify_api_key)])
async def segment_tarot(
    request: Request,
    file: Optional[UploadFile] = File(None),
    image_url: Optional[str] = Form(None),
    expected_count: int = Form(3)
):
    """
    塔罗牌分割专用接口
    1. 检测是否包含指定数量的塔罗牌 (默认为 3)
    2. 如果是，分别抠出这些牌并返回
    """
    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)
        # 固定 Prompt 检测塔罗牌
        output = processor.set_text_prompt(state=inference_state, prompt="tarot card")
        masks, boxes, scores = output["masks"], output["boxes"], output["scores"]
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"模型推理错误: {str(e)}")

    # 核心逻辑：判断数量
    detected_count = len(masks)
    
    # 创建本次请求的独立文件夹 (时间戳_UUID前8位)
    request_id = f"{int(time.time())}_{uuid.uuid4().hex[:8]}"
    output_dir = os.path.join(RESULT_IMAGE_DIR, request_id)
    os.makedirs(output_dir, exist_ok=True)

    if detected_count != expected_count:
        # 保存一张图用于调试/反馈
        try:
            filename = generate_and_save_result(image, inference_state, output_dir=output_dir)
            file_url = request.url_for("static", path=f"results/{request_id}/{filename}")
        except:
            file_url = None
        
        return JSONResponse(
            status_code=400, 
            content={
                "status": "failed",
                "message": f"检测到 {detected_count} 个目标，需要严格的 {expected_count} 张塔罗牌。请调整拍摄角度或背景。",
                "detected_count": detected_count,
                "debug_image_url": str(file_url) if file_url else None
            }
        )

    # 数量正确，执行抠图
    try:
        filenames = crop_and_save_objects(image, masks, boxes, output_dir=output_dir)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"抠图处理错误: {str(e)}")

    # 生成 URL 列表
    card_urls = [str(request.url_for("static", path=f"results/{request_id}/{fname}")) for fname in filenames]
    
    # 生成整体效果图
    try:
        main_filename = generate_and_save_result(image, inference_state, output_dir=output_dir)
        main_file_url = str(request.url_for("static", path=f"results/{request_id}/{main_filename}"))
    except:
        main_file_url = None

    return JSONResponse(content={
        "status": "success",
        "message": f"成功识别并分割 {expected_count} 张塔罗牌",
        "tarot_cards": card_urls,
        "full_visualization": main_file_url,
        "scores": scores.tolist() if torch.is_tensor(scores) else scores
    })

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