import os
import uuid
import yaml
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from ultralytics import YOLO

CONFIG_PATH = "./models/config.yaml"
WEIGHT_PATH = "./models/best.pt"
DATA_PATH = "./models/data.yaml"

# Load class names from data.yaml
with open(DATA_PATH, "r") as f:
    data_config = yaml.safe_load(f)
    CLASS_NAMES = data_config.get("names", [])

# Load YOLO model
model = YOLO(CONFIG_PATH, task="detect").load(WEIGHT_PATH)

app = FastAPI()

UPLOAD_FOLDER = "uploads"
OUTPUT_FOLDER = "output"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

@app.post("/predict")
async def predict(image: UploadFile = File(...)):
    filename = str(uuid.uuid4()) + "_" + image.filename
    img_path = os.path.join(UPLOAD_FOLDER, filename)
    
    with open(img_path, "wb") as buffer:
        buffer.write(await image.read())

    results = model(img_path, conf=0.2)

    detected_classes = set()

    for result in results:
        for box in result.boxes.data.tolist():
            _, _, _, _, _, class_id = box
            if 0 <= int(class_id) < len(CLASS_NAMES):
                detected_classes.add(CLASS_NAMES[int(class_id)])

    output_img_path = os.path.join(OUTPUT_FOLDER, filename)
    results[0].save(filename=output_img_path)

    os.remove(img_path)

    if not detected_classes:
        return JSONResponse(content={"error": "No objects detected"}, status_code=204)

    return JSONResponse(content={"items": list(detected_classes), "annotated_image": output_img_path})