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

# Setup logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

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

# Log model file presence
model_files = ["car_part_detector_model.pt", "damage_general_model.pt"]
for model_file in model_files:
    if os.path.exists(model_file):
        logger.info(f"Model file found: {model_file}")
    else:
        logger.error(f"Model file missing: {model_file}")

# Load YOLO models
try:
    logger.info("Loading car part model...")
    car_part_model = YOLO("car_part_detector_model.pt")
    logger.info("Car part model loaded successfully")
    logger.info("Loading damage model...")
    damage_model = YOLO("damage_general_model.pt")
    logger.info("Damage model loaded successfully")
except Exception as e:
    logger.error(f"Failed to load models: {str(e)}")
    raise RuntimeError(f"Failed to load models: {str(e)}")

def image_to_base64(img: np.ndarray) -> str:
    """Convert numpy image to base64 string."""
    try:
        _, buffer = cv2.imencode(".png", img)
        return base64.b64encode(buffer).decode("utf-8")
    except Exception as e:
        logger.error(f"Error encoding image to base64: {str(e)}")
        raise

@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."""
    logger.info("Received image upload")
    try:
        contents = await file.read()
        image = Image.open(BytesIO(contents)).convert("RGB")
        img = np.array(image)
        logger.info(f"Image loaded: shape={img.shape}")

        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()
        car_part_text = "Car Parts: No detections"
        damage_text = "Damage: No detections"

        try:
            logger.info("Running car part detection...")
            car_part_results = car_part_model(img)[0]
            if car_part_results.boxes:
                car_part_img = car_part_results.plot()[..., ::-1]
                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)
                )
            logger.info("Car part detection completed")
        except Exception as e:
            car_part_text = f"Car Parts: Error: {str(e)}"
            logger.error(f"Car part detection error: {str(e)}")

        try:
            logger.info("Running damage detection...")
            damage_results = damage_model(img)[0]
            if damage_results.boxes:
                damage_img = damage_results.plot()[..., ::-1]
                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)
                )
            logger.info("Damage detection completed")
        except Exception as e:
            damage_text = f"Damage: Error: {str(e)}"
            logger.error(f"Damage detection error: {str(e)}")

        car_part_img_base64 = image_to_base64(car_part_img)
        damage_img_base64 = image_to_base64(damage_img)
        logger.info("Returning prediction results")
        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:
        logger.error(f"Inference error: {str(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."""
    logger.info("Health check accessed")
    return {"message": "Car Parts & Damage Detection API is running"}