from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from ultralytics import YOLO
from PIL import Image
import numpy as np
import cv2
import io
import mediapipe as mp

app = FastAPI()

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

model = YOLO("app/pigmentnew.pt")
print("✅ YOLO Model Loaded. Classes:", model.names)

def detect_and_zoom_face(image_pil, zoom_scale=3.0):
    image_rgb = np.array(image_pil)
    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)

    with mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.6) as detector:
        results = detector.process(image_bgr)

        if not results.detections:
            return None, "No face detected"

        detection = results.detections[0]
        box = detection.location_data.relative_bounding_box
        h, w, _ = image_bgr.shape

        x = int(box.xmin * w)
        y = int(box.ymin * h)
        box_w = int(box.width * w)
        box_h = int(box.height * h)

        pad_x = int(box_w * 0.4)
        pad_y_top = int(box_h * 1.0)
        pad_y_bottom = int(box_h * 0.15)

        x1 = max(0, x - pad_x)
        y1 = max(0, y - pad_y_top)
        x2 = min(w, x + box_w + pad_x)
        y2 = min(h, y + box_h + pad_y_bottom)

        face_crop = image_rgb[y1:y2, x1:x2]
        zoomed = cv2.resize(face_crop, None, fx=zoom_scale, fy=zoom_scale, interpolation=cv2.INTER_CUBIC)
        zoomed_pil = Image.fromarray(zoomed)
        return zoomed_pil, None

@app.post("/zoom_face")
async def zoom_face(image: UploadFile = File(...)):
    try:
        contents = await image.read()
        image_pil = Image.open(io.BytesIO(contents)).convert('RGB')
        zoomed_face, error = detect_and_zoom_face(image_pil)
        if error:
            return JSONResponse(content={"error": error}, status_code=400)
        buf = io.BytesIO()
        zoomed_face.save(buf, format='JPEG')
        buf.seek(0)
        return StreamingResponse(buf, media_type="image/jpeg")
    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)

@app.post("/predict")
async def predict(image: UploadFile = File(...)):
    if not image:
        raise HTTPException(status_code=400, detail="No image provided")
    try:
        contents = await image.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid image: {str(e)}")

    results = model.predict(img, conf=0.09, imgsz=1024)
    result = results[0]
    detections = []
    if result.boxes is not None:
        for box in result.boxes:
            cls_id = int(box.cls[0].item())
            label = result.names[cls_id]
            confidence = round(float(box.conf[0].item()), 2)
            coords = box.xyxy[0].tolist()
            box_coords = [round(c, 2) for c in coords]
            detections.append({"label": label, "confidence": confidence, "box": box_coords})
    return {"detections": detections}