import logging
from fastapi import FastAPI, UploadFile, File, HTTPException, Form
from fastapi.responses import StreamingResponse
from rembg import remove
from PIL import Image
import io
import uvicorn
import asyncio
import numpy as np
from fastapi.middleware.cors import CORSMiddleware

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger("visa-backend")

app = FastAPI(title="Visa Photo Maker API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
def read_root():
    return {"status": "ok", "message": "Visa Photo Maker Backend is running"}

def add_white_background(
    image: Image.Image, 
    size: tuple = (600, 600),
    crop_params: dict = None
) -> Image.Image:
    """
    1. 裁剪掉透明边缘 (如果没有提供手动裁剪参数)
    2. 创建白色背景
    3. 将人物主体缩放并居中
    """
    # 确保图片是 RGBA 模式
    image = image.convert("RGBA")
    logger.info(f"Original image size: {image.size}, Target size: {size}")
    
    if crop_params and all(k in crop_params for k in ('x', 'y', 'w', 'h')):
        # --- 使用手动裁剪参数 ---
        logger.info(f"Using manual crop params: {crop_params}")
        x, y, w, h = crop_params['x'], crop_params['y'], crop_params['w'], crop_params['h']
        # 裁剪原图 (crop_params 应该是相对于原图的像素坐标)
        image = image.crop((x, y, x + w, y + h))
    else:
        # --- 自动强力裁剪逻辑开始 ---
        # 将 PIL Image 转为 numpy 数组以处理 Alpha 通道
        img_np = np.array(image)
        
        # 获取 Alpha 通道 (第4个通道)
        alpha = img_np[:, :, 3]
        
        # 设定阈值：Alpha 值小于 50 的像素视为完全透明
        threshold = 50
        mask = alpha > threshold
        
        # 如果全图都是透明的（没有检测到人），就直接返回白图
        if not np.any(mask):
            logger.warning("No subject found (image is empty)")
            return Image.new("RGB", size, (255, 255, 255))
            
        # 获取非零区域的坐标 (行, 列)
        rows = np.any(mask, axis=1)
        cols = np.any(mask, axis=0)
        
        y_min, y_max = np.where(rows)[0][[0, -1]]
        x_min, x_max = np.where(cols)[0][[0, -1]]
        
        # 裁剪图片
        image = image.crop((x_min, y_min, x_max + 1, y_max + 1))
        logger.info(f"Cropped size (tight): {image.size}")
        
        # --- 智能半身裁剪 (Smart Body Crop) ---
        target_ratio = size[1] / size[0]
        max_allowed_ratio = target_ratio + 0.15
        
        w, h = image.size
        current_ratio = h / w
        
        if current_ratio > max_allowed_ratio:
            logger.info(f"Image is too tall (ratio {current_ratio:.2f} > {max_allowed_ratio:.2f}), cropping bottom...")
            new_h = int(w * max_allowed_ratio)
            image = image.crop((0, 0, w, new_h))
            logger.info(f"New size after smart crop: {image.size}")
    # ------------------------------------

    # 3. 计算缩放比例：Cover 模式，但稍微留一点边距 (95% 覆盖)
    scale_width = (size[0] * 0.95) / image.width
    scale_height = (size[1] * 0.95) / image.height
    
    scale_factor = max(scale_width, scale_height)
    
    new_width = int(image.width * scale_factor)
    new_height = int(image.height * scale_factor)
    
    logger.info(f"Resizing: {new_width}x{new_height}")
    resized_img = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    final_image = Image.new("RGB", size, (255, 255, 255))
    x_offset = (size[0] - new_width) // 2
    
    top_padding = int(size[1] * 0.10)
    y_offset = top_padding
    
    if y_offset + new_height < size[1]:
        y_offset = (size[1] - new_height) // 2
    elif y_offset + new_height > size[1] + 100:
        y_offset = int(size[1] * 0.05)

    final_image.paste(resized_img, (x_offset, y_offset), resized_img)
    
    return final_image.convert("RGB")

# Global lock to prevent CPU overload
processing_lock = asyncio.Lock()

@app.post("/process-image")
async def process_image_endpoint(
    file: UploadFile = File(...),
    width: int = Form(600),
    height: int = Form(600),
    crop_x: int = Form(None),
    crop_y: int = Form(None),
    crop_w: int = Form(None),
    crop_h: int = Form(None)
):
    # Wait for the lock before processing
    async with processing_lock:
        try:
            input_image_bytes = await file.read()
            logger.info(f"Processing image for file: {file.filename}")
            
            output_image_bytes = await asyncio.to_thread(remove, input_image_bytes)
            img_no_bg = Image.open(io.BytesIO(output_image_bytes))
            
            crop_params = None
            if crop_x is not None and crop_y is not None:
                crop_params = {'x': crop_x, 'y': crop_y, 'w': crop_w, 'h': crop_h}
            
            final_image = add_white_background(img_no_bg, size=(width, height), crop_params=crop_params)
            
            img_io = io.BytesIO()
            final_image.save(img_io, format="JPEG", quality=95)
            img_io.seek(0)
            
            logger.info(f"Successfully processed image: {file.filename}")
            return StreamingResponse(img_io, media_type="image/jpeg")
            
        except Exception as e:
            logger.error(f"Error processing image: {str(e)}", exc_info=True)
            raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=13002)