from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from ultralytics import YOLO
import numpy as np
import cv2
from io import BytesIO
from PIL import Image
import base64

app = FastAPI(title="Car Parts & Damage Detection API")

# Load YOLO models
try:
    car_part_model = YOLO("car_part_detector_model.pt")
    damage_model = YOLO("damage_general_model.pt")
except Exception as e:
    raise RuntimeError(f"Failed to load models: {str(e)}")

def image_to_base64(img: np.ndarray) -> str:
    """Convert numpy image to base64 string for JSON response."""
    _, buffer = cv2.imencode(".png", img)
    return base64.b64encode(buffer).decode("utf-8")

@app.post("/predict", summary="Run inference on an image for car parts and damage")
async def predict(file: UploadFile = File(...)):
    """
    Upload an image and get car parts and damage detection results.
    Returns annotated images as base64 strings and text descriptions.
    """
    try:
        # Read and process image
        contents = await file.read()
        image = Image.open(BytesIO(contents)).convert("RGB")
        img = np.array(image)

        # Initialize default blank images (gray placeholder)
        blank_img = np.full((img.shape[0], img.shape[1], 3), 128, dtype=np.uint8)
        car_part_img = blank_img.copy()
        damage_img = blank_img.copy()

        # Initialize text results
        car_part_text = "Car Parts: No detections"
        damage_text = "Damage: No detections"

        # Process car parts detection
        try:
            car_part_results = car_part_model(img)[0]
            if car_part_results.boxes:
                car_part_img = car_part_results.plot()[..., ::-1]  # BGR to RGB
                car_part_text = "Car Parts:\n" + "\n".join(
                    f"- {car_part_results.names[int(cls)]} ({conf:.2f})"
                    for conf, cls in zip(car_part_results.boxes.conf, car_part_results.boxes.cls)
                )
        except Exception as e:
            car_part_text = f"Car Parts: Error: {str(e)}"

        # Process damage detection
        try:
            damage_results = damage_model(img)[0]
            if damage_results.boxes:
                damage_img = damage_results.plot()[..., ::-1]  # BGR to RGB
                damage_text = "Damage:\n" + "\n".join(
                    f"- {damage_results.names[int(cls)]} ({conf:.2f})"
                    for conf, cls in zip(damage_results.boxes.conf, damage_results.boxes.cls)
                )
        except Exception as e:
            damage_text = f"Damage: Error: {str(e)}"

        # Convert output images to base64
        car_part_img_base64 = image_to_base64(car_part_img)
        damage_img_base64 = image_to_base64(damage_img)

        return JSONResponse({
            "car_part_image": car_part_img_base64,
            "car_part_text": car_part_text,
            "damage_image": damage_img_base64,
            "damage_text": damage_text
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")

@app.get("/", summary="Health check")
async def root():
    """Check if the API is running."""
    return {"message": "Car Parts & Damage Detection API is running"}