from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
from PIL import Image
import io
import os

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Download and cache model on startup
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(
    repo_id="Decizez/yolov-corrosion-detection",
    filename="Lite_YOLO8_v1.pt"
)

model = YOLO(model_path)
print("✅ Model loaded successfully")

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

@app.post("/detect")
async def detect(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert("RGB")
    w, h = image.width, image.height

    results = model(image)

    detections = []
    for result in results:
        boxes = result.boxes
        if boxes is not None:
            for box in boxes:
                conf = float(box.conf[0])
                if conf > 0.3:
                    x1, y1, x2, y2 = box.xyxy[0].tolist()
                    detections.append({
                        "label": "Corrosion Detected",
                        "confidence": round(conf * 100, 1),
                        "area_percent": round(
                            ((x2 - x1) * (y2 - y1)) / (w * h) * 100, 1
                        ),
                        "box": {
                            "x": round(x1 / w * 100, 1),
                            "y": round(y1 / h * 100, 1),
                            "width": round((x2 - x1) / w * 100, 1),
                            "height": round((y2 - y1) / h * 100, 1)
                        }
                    })

    return {
        "detections": detections,
        "frame_size": {"w": w, "h": h}
    }