from fastapi import FastAPI, File, UploadFile
from ultralytics import YOLO
import shutil
import uuid
from pathlib import Path
import base64

# Init FastAPI
app = FastAPI()

# Load YOLO model (place best.pt in /app folder)
model = YOLO("best.pt")

@app.get("/")
async def root():
    return {"status": "ok", "message": "YOLOv8 Yellow Disease Detection API is running 🚀"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    temp_file = Path(f"/tmp/{uuid.uuid4()}_{file.filename}")
    with temp_file.open("wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    results = model.predict(source=str(temp_file), save=False, conf=0.5)

    detections = []
    disease_found = False

    for r in results:
        for box in r.boxes:
            cls_name = model.names[int(box.cls)]
            conf = float(box.conf)
            detections.append({
                "class": cls_name,
                "confidence": conf,
                "bbox": [float(v) for v in box.xyxy[0].tolist()]
            })
            if cls_name == "yellow_disease":
                disease_found = True

    temp_file.unlink(missing_ok=True)

    return {
        "spray_signal": 1 if disease_found else 0,   # ESC reads this
        "frame_classification": "yellow_disease" if disease_found else "healthy",
        "detections": detections                     # keep for logs/debugging
    }