import os


DAMAGE_CHECK_PARTS = {
    'driver_front_side',
    'driver_rear_side',
    'passenger_front_side',
    'passenger_rear_side',
}


def _class_name(model, class_index):
    """Read a class label from YOLO names whether it is a dict or a list."""
    names = model.names
    if isinstance(names, dict):
        return names.get(class_index, 'unknown')
    if 0 <= class_index < len(names):
        return names[class_index]
    return 'unknown'


def run_single_inference(model, filepath):
    """Run one YOLO model on one image and normalize the top prediction."""
    if model is None:
        raise RuntimeError('Inference model is not loaded.')

    results = model(filepath, verbose=False)
    if not results:
        return {'class': 'unknown', 'confidence': 0.0}

    result = results[0]

    if result.probs is not None:
        class_index = int(result.probs.top1)
        confidence = float(result.probs.top1conf)
        class_name = _class_name(model, class_index)
    elif result.boxes is not None and len(result.boxes) > 0:
        class_index = int(result.boxes.cls[0].item())
        confidence = float(result.boxes.conf[0].item())
        class_name = _class_name(model, class_index)
    else:
        class_name = 'unknown'
        confidence = 0.0

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


def process_images(parts_model, damage_model, image_inputs):
    """
    Process uploaded images while preserving each browser-visible filename/index.

    image_inputs accepts dictionaries with path, filename, and index. Plain path
    strings are also accepted for compatibility with older callers.
    """
    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 fallback_index, image_input in enumerate(image_inputs):
        if isinstance(image_input, dict):
            filepath = image_input['path']
            filename = image_input.get('filename') or os.path.basename(filepath)
            client_index = image_input.get('index', fallback_index)
        else:
            filepath = image_input
            filename = os.path.basename(filepath)
            client_index = fallback_index

        print(f'Processing {filename}...')

        part_prediction = run_single_inference(parts_model, filepath)
        predicted_part = part_prediction['class']

        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'.")
            damage_prediction = {
                'class': 'correct',
                'confidence': 1.0,
                'note': 'Result by default, not by model inference.',
            }

        final_results.append({
            'index': client_index,
            'filename': filename,
            'part_prediction': part_prediction,
            'damage_prediction': damage_prediction,
        })

    return final_results