import io
import time
import torch
from PIL import Image
from ultralytics import YOLO
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager

# ─── أسماء الـ 17 Keypoints ───────────────────────────────────────
KEYPOINT_NAMES = [
    "nose", "left_eye", "right_eye", "left_ear", "right_ear",
    "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
    "left_wrist", "right_wrist", "left_hip", "right_hip",
    "left_knee", "right_knee", "left_ankle", "right_ankle"
]

MODEL_DATA = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("📥 Loading YOLO11x-pose...")
    start = time.time()
    # يتحمل تلقائياً من ultralytics
    MODEL_DATA["model"] = YOLO("yolo11x-pose.pt")
    print(f"✅ YOLO11x-pose ready in {time.time()-start:.1f}s")
    yield
    MODEL_DATA.clear()

app = FastAPI(
    title="Human Body Detection - YOLO11x-pose",
    description="Detects human body and keypoints using YOLO11x-pose",
    version="1.0.0",
    lifespan=lifespan
)

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

@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": "model" in MODEL_DATA}

@app.post("/analyze")
async def analyze_image(file: UploadFile = File(...)):

    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="الملف ليس صورة")

    try:
        image = Image.open(io.BytesIO(await file.read())).convert("RGB")
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"خطأ في قراءة الصورة: {str(e)}")

    try:
        model = MODEL_DATA["model"]
        start_time = time.time()

        # ─── تشغيل YOLO ──────────────────────────────────────────
        results = model(image, verbose=False)
        elapsed = round(time.time() - start_time, 2)

        persons = []

        for r in results:
            num_persons = len(r.boxes) if r.boxes is not None else 0

            for i in range(num_persons):
                person = {"id": i + 1}

                # ─── Bounding Box ─────────────────────────────────
                if r.boxes is not None and i < len(r.boxes):
                    box = r.boxes.xyxy[i].tolist()
                    conf = float(r.boxes.conf[i])
                    person["bbox"]       = [round(x, 1) for x in box]
                    person["confidence"] = round(conf, 3)

                # ─── Keypoints ────────────────────────────────────
                if r.keypoints is not None and i < len(r.keypoints.xy):
                    kp_xy   = r.keypoints.xy[i].tolist()
                    kp_conf = r.keypoints.conf[i].tolist() if r.keypoints.conf is not None else [1.0] * 17

                    visible_keypoints = {}
                    for name, (x, y), c in zip(KEYPOINT_NAMES, kp_xy, kp_conf):
                        if x > 0 and y > 0 and c > 0.3:
                            visible_keypoints[name] = {
                                "x": round(x, 1),
                                "y": round(y, 1),
                                "confidence": round(c, 3)
                            }

                    person["visible_keypoints"]       = visible_keypoints
                    person["visible_keypoints_count"] = len(visible_keypoints)
                    person["visible_keypoints_names"] = list(visible_keypoints.keys())

                persons.append(person)

        # ─── القرار النهائي ───────────────────────────────────────
        human_detected = len(persons) > 0
        summary = f"yes detected human: {len(persons)} person(s)" if human_detected else "no detected human body"

        return {
            "summary":        summary,
            "detected":       human_detected,
            "persons_count":  len(persons),
            "persons":        persons,
            "execution_time": elapsed,
            "status":         "success"
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


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