from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn
import base64
import cv2
import numpy as np
from ultralytics import YOLO
from datetime import datetime

app = FastAPI(title="PCB Defect Detection API")
model = YOLO("Best_defects.pt")

class ImageRequest(BaseModel):
    """Request model for image processing endpoint."""
    image: str

@app.get("/")
async def root():
    """Root endpoint to verify API status."""
    current_time = datetime.now().isoformat()
    return {
        "message": "PCB Defects API works",
        "time": current_time
    }

@app.post("/predict")
async def predict(request: ImageRequest):
    """
    Process an image to detect PCB defects.
    
    Args:
        request: Contains base64 encoded image
        
    Returns:
        JSON with detection statistics and bounding boxes
    """
    # Validate image input
    if not request.image:
        return {"error": "Invalid Image"}
    
    try:
        image_bytes = base64.b64decode(request.image, validate=True)
        
        np_arr = np.frombuffer(image_bytes, np.uint8)
        image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        
        if image is None:
            return {"error": "Invalid image"}
            
        results = model.predict(image, conf=0.36)
        result = results[0]

        json_result = {}
        class_counters = {}
        
        for box in result.boxes:
            class_id = int(box.cls[0])
            class_name = result.names[class_id]
            x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
            
            if class_name in class_counters:
                class_counters[class_name] += 1
            else:
                class_counters[class_name] = 1
                
            key = f"{class_name}{class_counters[class_name]}" if class_counters[class_name] > 1 else class_name
            json_result[key] = [x1, y1, x2, y2]
        
        if hasattr(result, "summary") and isinstance(result.summary, dict):
            statistics_summary = result.summary
        else:
            statistics_summary = {name: count for name, count in class_counters.items()}
            
        return {
            "statistics": statistics_summary,
            "detections": json_result
        }
        
    except Exception as e:
        return {"error": f"Invalid Image: {str(e)}"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)