# processing.py

import os
from ultralytics import YOLO

# --- Configuration ---
# These are the specific parts that require a subsequent damage check.
DAMAGE_CHECK_PARTS = {
    'driver_front_side',
    'driver_rear_side',
    'passenger_front_side',
    'passenger_rear_side',
}

def run_single_inference(model, filepath):
    """
    Helper function to run inference for a single model and format the result.
    """
    if model is None:
        return None  # Return None if the model isn't loaded

    results = model(filepath, verbose=False) # verbose=False to keep logs clean
    result = results[0]
    
    # Check if it's a classification model with probabilities
    if result.probs is not None:
        probs = result.probs
        top1_index = probs.top1
        top1_confidence = float(probs.top1conf)
        class_name = model.names[top1_index]
    else: # Fallback for detection models or if probs are not available
        # Assuming the top prediction is what we need
        top1_index = result.boxes.cls[0].int() if len(result.boxes) > 0 else 0
        top1_confidence = float(result.boxes.conf[0]) if len(result.boxes) > 0 else 0.0
        class_name = model.names[top1_index] if len(result.boxes) > 0 else "unknown"

    return {
        "class": class_name,
        "confidence": round(top1_confidence, 4)
    }

def process_images(parts_model, damage_model, image_paths):
    """
    Processes a list of images.
    1. Runs the 'parts_model' on every image.
    2. If the detected part is in DAMAGE_CHECK_PARTS, it then runs the 'damage_model'.
    3. Otherwise, the damage status defaults to 'correct'.
    """
    if parts_model is None or damage_model is None:
        raise RuntimeError("One or more models are not loaded. Check server logs.")

    final_results = []

    for filepath in image_paths:
        filename = os.path.basename(filepath)
        print(f"Processing {filename}...")

        # 1. First, predict the part
        part_prediction = run_single_inference(parts_model, filepath)
        predicted_part = part_prediction.get("class") if part_prediction else "unknown"

        damage_prediction = None
        # 2. Conditionally predict the damage
        if predicted_part in DAMAGE_CHECK_PARTS:
            print(f"  -> Part '{predicted_part}' requires damage check. Running damage model...")
            damage_prediction = run_single_inference(damage_model, filepath)
        else:
            print(f"  -> Part '{predicted_part}' does not require damage check. Defaulting to 'correct'.")
            # 3. For other parts, default to 'correct'
            damage_prediction = {
                "class": "correct",
                "confidence": 1.0,
                "note": "Result by default, not by model inference."
            }

        # Assemble the final result for this image
        final_results.append({
            "filename": filename,
            "part_prediction": part_prediction,
            "damage_prediction": damage_prediction
        })

    return final_results