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

# --- DEBUGGING STEP 1 ---
print("Application script is starting...")

os.environ['TORCH_HOME'] = '/tmp/torch_cache'

app = FastAPI(title="YOLOv8 Car Damage Detection API")

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

model = YOLO('best.pt')


@app.get("/")
def read_root():
    return {"message": "Welcome to the Car Damage Detection API!"}


@app.post("/detect")
async def detect_damage(file: UploadFile = File(...)):
    # 1. Get image and run the model (same as before)
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert("RGB")
    results = model(image)

    # 2. Process the results to determine unique conditions

    # Use a set to automatically handle duplicates (e.g., "broken" is only added once)
    conditions = set()

    # Define which detections count as "broken" or "dirty"
    broken_classes = {'scratch', 'rust', 'dent'}
    dirty_classes = {'dirt'}

    # Loop through all detected objects
    for result in results:
        for box in result.boxes:
            # Get the name and confidence of the detected object
            class_name = model.names[int(box.cls[0])]
            confidence = float(box.conf[0])

            # Only consider detections with high confidence
            if confidence > 0.5:
                if class_name in broken_classes:
                    conditions.add("Broken")
                elif class_name in dirty_classes:
                    conditions.add("Dirty")

    # 3. Prepare the final output

    # If the conditions set is empty after checking all boxes, the car is clean
    if not conditions:
        final_output = ["No damage detected."]
    else:
        # Convert the set to a list for the JSON response
        final_output = list(conditions)

    # Return the simple list of conditions
    return {"conditions": final_output}


# --- DEBUGGING STEP 2 ---
# This will print all routes that FastAPI has successfully registered.
print("\n--- Registered Routes ---")
for route in app.routes:
    print(f"Path: {route.path}, Methods: {', '.join(route.methods)}")
print("-------------------------\n")