from fastapi import FastAPI, UploadFile, File, HTTPException, Security, Depends
from fastapi.security.api_key import APIKeyHeader
from fastapi.responses import JSONResponse, StreamingResponse
import uvicorn
import logging
import io
import time
import numpy as np
from PIL import Image
import cv2
from ultralytics import YOLO
import mediapipe as mp

# ==========================
# 🔑 Sécurité : API Key
# ==========================
API_KEY = "1234"  # ⚠️ change avant de partager
api_key_header = APIKeyHeader(name="X-API-Key")

def verify_api_key(api_key: str = Security(api_key_header)):
    if api_key != API_KEY:
        raise HTTPException(status_code=403, detail="Forbidden")
    return api_key

# ==========================
# 📝 Logger
# ==========================
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("stroke-api")

# ==========================
# 🚀 App
# ==========================
app = FastAPI(
    title="Stroke Detection API",
    version="1.2.0",
    description="🚑 Stroke Detection API using YOLOv8 + Face Detection (MediaPipe). Research/demo only."
)

# ==========================
# 📦 Chargement modèles
# ==========================
try:
    model = YOLO("best.pt")
    logger.info("✅ YOLO model loaded.")
except Exception as e:
    logger.exception("❌ Failed to load YOLO model")
    raise RuntimeError(f"Model loading failed: {e}")

mp_face_detection = mp.solutions.face_detection

# ==========================
# 🔧 Utilitaires
# ==========================
ALLOWED_EXT = (".png", ".jpg", ".jpeg")
ALLOWED_MIME = {"image/png", "image/jpeg"}
MAX_BYTES = 8 * 1024 * 1024  # 8 MB
CROP_ON_FACE = True

def _validate_file(file: UploadFile, raw: bytes):
    if not file.filename.lower().endswith(ALLOWED_EXT):
        raise HTTPException(status_code=400, detail="Invalid file extension")
    if (file.content_type or "").lower() not in ALLOWED_MIME and file.content_type:
        raise HTTPException(status_code=400, detail="Invalid content-type")
    if len(raw) > MAX_BYTES:
        raise HTTPException(status_code=413, detail=f"Image too large. Max {MAX_BYTES//(1024*1024)} MB")

def _read_image_to_numpy(raw: bytes) -> np.ndarray:
    try:
        img = Image.open(io.BytesIO(raw)).convert("RGB")
        return np.array(img)
    except Exception:
        raise HTTPException(status_code=400, detail="Unreadable image file")

def _largest_face_bbox(np_img: np.ndarray, min_conf: float = 0.6):
    h, w = np_img.shape[:2]
    with mp_face_detection.FaceDetection(min_detection_confidence=min_conf) as fd:
        results = fd.process(cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR))
        if not results.detections:
            return None
        boxes = []
        for det in results.detections:
            rel = det.location_data.relative_bounding_box
            x1 = int(max(0, rel.xmin) * w)
            y1 = int(max(0, rel.ymin) * h)
            x2 = int(min(1.0, rel.xmin + rel.width) * w)
            y2 = int(min(1.0, rel.ymin + rel.height) * h)
            boxes.append((x1, y1, x2, y2))
        boxes.sort(key=lambda b: (b[2]-b[0])*(b[3]-b[1]), reverse=True)
        return boxes[0] if boxes else None

def _crop_to_bbox(np_img: np.ndarray, bbox, margin: float = 0.15) -> np.ndarray:
    h, w = np_img.shape[:2]
    x1, y1, x2, y2 = bbox
    bw, bh = x2 - x1, y2 - y1
    dx, dy = int(bw*margin), int(bh*margin)
    X1, Y1 = max(0,x1-dx), max(0,y1-dy)
    X2, Y2 = min(w,x2+dx), min(h,y2+dy)
    return np_img[Y1:Y2, X1:X2].copy()

def _annotate_face_box(np_img: np.ndarray, bbox) -> np.ndarray:
    annotated = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR).copy()
    x1, y1, x2, y2 = bbox
    cv2.rectangle(annotated, (x1, y1), (x2, y2), (0,255,0), 2)
    return annotated

# ==========================
# 🩺 Healthcheck
# ==========================
@app.get("/health")
async def health():
    return {"status": "ok", "model_loaded": True}

# ==========================
# 📦 Endpoint JSON
# ==========================
@app.post("/v1/predict/")
async def predict(file: UploadFile = File(...), api_key: str = Depends(verify_api_key)):
    raw = await file.read()
    _validate_file(file, raw)
    try:
        np_img = _read_image_to_numpy(raw)
        face_bbox = _largest_face_bbox(np_img)
        if face_bbox is None:
            return JSONResponse(status_code=422, content={"status":"error","message":"Aucun visage détecté"})
        input_img = _crop_to_bbox(np_img, face_bbox) if CROP_ON_FACE else np_img

        start_time = time.time()
        results = model.predict(source=input_img, verbose=False)
        elapsed = time.time() - start_time

        output = [{"class": r.names[int(box.cls[0].item())],
                   "confidence": float(box.conf[0].item()),
                   "bbox": box.xyxy[0].tolist()} 
                  for r in results for box in r.boxes]

        return JSONResponse(content={
            "status": "ok",
            "face_detected": True,
            "face_bbox": list(map(int, face_bbox)),
            "predictions": output,
            "inference_time_sec": round(elapsed,3)
        })

    except Exception as e:
        logger.exception("Error in /v1/predict")
        raise HTTPException(status_code=500, detail=str(e))

# ==========================
# 🖼️ Endpoint Image annotée
# ==========================
@app.post("/v1/predict_image/")
async def predict_image(file: UploadFile = File(...), api_key: str = Depends(verify_api_key)):
    raw = await file.read()
    _validate_file(file, raw)
    try:
        np_img = _read_image_to_numpy(raw)
        face_bbox = _largest_face_bbox(np_img)
        if face_bbox is None:
            return JSONResponse(status_code=422, content={"status":"error","message":"Aucun visage détecté"})
        input_img = _crop_to_bbox(np_img, face_bbox) if CROP_ON_FACE else np_img

        start_time = time.time()
        results = model.predict(source=input_img, verbose=False)
        elapsed = time.time() - start_time

        yolo_annot = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
        out_rgb = yolo_annot if CROP_ON_FACE else _annotate_face_box(np_img, face_bbox)

        pil_img = Image.fromarray(out_rgb)
        buf = io.BytesIO()
        pil_img.save(buf, format="PNG")
        buf.seek(0)

        headers = {"X-Inference-Time": str(round(elapsed,3))}
        return StreamingResponse(buf, media_type="image/png", headers=headers)

    except Exception as e:
        logger.exception("Error in /v1/predict_image")
        raise HTTPException(status_code=500, detail=str(e))

# ==========================
# 🧪 Test automatique
# ==========================
@app.get("/test_upload/")
async def test_upload():
    try:
        file_path = "test.jpg"
        np_img = _read_image_to_numpy(open(file_path,"rb").read())
        face_bbox = _largest_face_bbox(np_img)
        if not face_bbox:
            return {"status":"error","message":"Aucun visage détecté"}
        results = model.predict(source=np_img, verbose=False)
        return {"status":"ok","face_detected":True,"num_detections":len(results[0].boxes)}
    except Exception as e:
        return {"status":"error","message": str(e)}

# ==========================
# 🚀 Lancement local
# ==========================
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)