import gradio as gr
import numpy as np
from ultralytics import YOLO
import insightface
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import cv2

# ----------------------------
# Load models (CPU)
# ----------------------------

yolo = YOLO("yolov8n.pt")

face_model = insightface.app.FaceAnalysis(name="buffalo_l")
face_model.prepare(ctx_id=-1)


def normalize(vec):
    vec = np.array(vec, dtype=np.float32)
    norm = np.linalg.norm(vec)
    if norm == 0:
        return vec.tolist()
    return (vec / norm).tolist()


def process_image_np(image_np):
    results = yolo(image_np)
    faces_output = []

    for r in results:
        boxes = r.boxes
        for box, cls, conf in zip(boxes.xyxy, boxes.cls, boxes.conf):
            if int(cls) != 0:
                continue
            if float(conf) < 0.4:
                continue

            xmin, ymin, xmax, ymax = box.cpu().numpy()
            xmin, ymin, xmax, ymax = map(int, [xmin, ymin, xmax, ymax])

            h, w, _ = image_np.shape
            xmin = max(0, xmin)
            ymin = max(0, ymin)
            xmax = min(w, xmax)
            ymax = min(h, ymax)

            person_crop = image_np[ymin:ymax, xmin:xmax]
            if person_crop.size == 0:
                continue

            detected_faces = face_model.get(person_crop)

            for face in detected_faces:
                embedding = normalize(face.embedding)
                fxmin, fymin, fxmax, fymax = face.bbox.astype(int)

                faces_output.append({
                    "cx": float((fxmin + fxmax) / 2 + xmin),
                    "cy": float((fymin + fymax) / 2 + ymin),
                    "confidence": float(conf),
                    "box": {
                        "xmin": int(fxmin + xmin),
                        "ymin": int(fymin + ymin),
                        "xmax": int(fxmax + xmin),
                        "ymax": int(fymax + ymin)
                    },
                    "embedding": embedding
                })

    return faces_output


# ----------------------------------------
# Create Gradio Blocks
# ----------------------------------------

with gr.Blocks() as demo:
    gr.Markdown("Face Processing API Running")

# Get underlying FastAPI app from Gradio
app = demo.app


# ----------------------------------------
# ADD REAL REST ENDPOINT TO SAME APP
# ----------------------------------------

@app.post("/detect")
async def detect(request: Request):
    body = await request.body()

    np_arr = np.frombuffer(body, np.uint8)
    image_np = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

    if image_np is None:
        return JSONResponse({"error": "Invalid image"}, status_code=400)

    result = process_image_np(image_np)
    return result


# Launch
demo.queue(False)
demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)