from fastapi import FastAPI, File, UploadFile
from fastapi.responses import Response, JSONResponse
from insightface.app import FaceAnalysis
from rembg import remove
from PIL import Image
import numpy as np
import io

app = FastAPI()

# Simple homepage (fixes 404 on HF Spaces)
@app.get("/")
def home():
    return {"message": "ID Photo API is running. Use POST /process"}

# Load face detection model
face_app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
face_app.prepare(ctx_id=0, det_size=(640, 640))


def resize_to_4x6(img):
    return img.resize((472, 709), Image.LANCZOS)


@app.post("/process")
async def process_img(file: UploadFile = File(...)):
    try:
        img_bytes = await file.read()
        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
        np_img = np.array(img)

        faces = face_app.get(np_img)
        if not faces:
            return JSONResponse({"error": "No face detected"}, status_code=400)

        face = faces[0]

        # --- Extract face bounding box ---
        x1, y1, x2, y2 = face.bbox.astype(int)
        cropped = img.crop((x1, y1, x2, y2))

        # --- Resize cropped face with preserved aspect ratio ---
        max_face_height = int(709 * 0.75)  # face occupies 75% of final height
        w, h = cropped.size
        scale_factor = max_face_height / h

        new_w = int(w * scale_factor)
        new_h = int(h * scale_factor)
        resized_face = cropped.resize((new_w, new_h), Image.LANCZOS)

        # --- Create final white 4x6 canvas ---
        final_w, final_h = 472, 709
        canvas = Image.new("RGB", (final_w, final_h), (255, 255, 255))

        # --- Center image on canvas ---
        paste_x = (final_w - new_w) // 2
        paste_y = (final_h - new_h) // 2
        canvas.paste(resized_face, (paste_x, paste_y))

        # --- Output final image ---
        buf = io.BytesIO()
        canvas.save(buf, format="JPEG")
        buf.seek(0)
        return Response(buf.getvalue(), media_type="image/jpeg")

    except Exception as e:
        print("ERROR:", e)
        return JSONResponse({"error": str(e)}, status_code=500)