import io
import os
from typing import List

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

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

_model = None
MODEL_CACHE_DIR = os.environ.get("MODEL_CACHE_DIR", "/models")


def get_model():
    global _model
    if _model is None:
        from ultralytics import YOLO
        os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
        _model = YOLO("foduucom/plant-leaf-detection-and-classification")
    return _model


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/detect")
async def detect(file: UploadFile = File(...)):
    data = await file.read()
    img = Image.open(io.BytesIO(data)).convert("RGB")

    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