from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import cv2
import numpy as np
import base64
from ultralytics import YOLO

app = FastAPI()

# Enable CORS so your React frontend/Node backend can query this Space
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load lightweight YOLOv8 Nano model (pretrained on COCO dataset)
model = YOLO("yolov8n.pt")

class ImagePayload(BaseModel):
    image: str

@app.get("/")
def home():
    return {"status": "YOLOv8 Active", "model": "yolov8n"}

@app.post("/predict")
def predict(payload: ImagePayload):
    try:
        # Decode base64 image
        encoded_data = payload.image.split(',')[1] if ',' in payload.image else payload.image
        nparr = np.frombuffer(base64.b64decode(encoded_data), np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if img is None:
            return {"phoneDetected": False, "error": "Invalid image data"}

        # Run inference
        results = model(img)
        phone_detected = False
        
        for r in results:
            for box in r.boxes:
                class_id = int(box.cls[0])
                label = model.names[class_id]
                # Label 'cell phone' or 'laptop' or 'remote' in COCO dataset
                if label in ['cell phone', 'laptop', 'remote']:
                    phone_detected = True
                    break
                    
        return {"phoneDetected": phone_detected}
    except Exception as e:
        return {"phoneDetected": False, "error": str(e)}