import asyncio
import io
import os
from typing import List

from fastapi import FastAPI, File, HTTPException, UploadFile
from PIL import Image, UnidentifiedImageError

app = FastAPI(title="AI Plant Detection Service")

_model = None
MODEL_CACHE_DIR = os.environ.get("MODEL_CACHE_DIR", "/models")
MAX_FILE_SIZE = 10 * 1024 * 1024  # 10 MB


def get_model():
    global _model
    if _model is None:
        from ultralytics import YOLO
        os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
        try:
            _model = YOLO("foduucom/plant-leaf-detection-and-classification")
        except Exception as e:
            raise RuntimeError(f"Impossible de charger le modèle YOLO: {e}") from e
    return _model


def _run_inference(img: Image.Image) -> list:
    model = get_model()
    results = model.predict(img, conf=0.25, iou=0.45, verbose=False)
    detections = []
    if results and results[0].boxes:
        boxes = results[0].boxes
        names = model.names
        for i in range(min(3, len(boxes))):
            cls_id = int(boxes.cls[i].item())
            conf = float(boxes.conf[i].item())
            detections.append({
                "class_name": names[cls_id],
                "confidence": round(conf, 3),
            })
    return detections


@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": _model is not None}


@app.post("/detect")
async def detect(file: UploadFile = File(...)):
    data = await file.read()

    if len(data) > MAX_FILE_SIZE:
        raise HTTPException(status_code=413, detail="Fichier trop volumineux (max 10 MB)")

    try:
        img = Image.open(io.BytesIO(data)).convert("RGB")
    except (UnidentifiedImageError, OSError) as e:
        raise HTTPException(status_code=400, detail=f"Image invalide: {e}")

    try:
        loop = asyncio.get_event_loop()
        detections = await loop.run_in_executor(None, _run_inference, img)
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))

    return detections