# Standard library imports
# (Add any necessary imports for future object detection implementation)
import json

# Third-party imports
from PIL import Image
import numpy as np
from ultralytics import YOLO

# Local imports
from utils.image_utils import load_image, preprocess_image


# Load the YOLOv8 model globally to avoid reloading on each function call
# Using a common pre-trained YOLOv8 nano model ('yolov8n.pt')
try:
    model = YOLO('yolov8n.pt')
    print("YOLOv8 model loaded successfully.")
except Exception as e:
    print(f"Error loading YOLOv8 model: {e}")
    model = None # Set model to None if loading fails


def object_detection(input_type, uploaded_image, image_url, base64_string):
    """
    Performs object detection on the image from various input types using YOLOv8.

    Args:
        input_type (str): The selected input method ("Upload File", "Enter URL", "Enter Base64").
        uploaded_image (PIL.Image.Image): The uploaded image (if input_type is "Upload File").
        image_url (str): The image URL (if input_type is "Enter URL").
        base64_string (str): The image base64 string (if input_type is "Enter Base64").

    Returns:
        tuple: A tuple containing:
            - numpy.ndarray: The image with detected objects drawn on it, or None if an error occurred or model not loaded.
            - dict: A dictionary containing the raw detection data (bounding boxes, classes, scores), or None.
    """
    if model is None:
        print("YOLOv8 model is not loaded. Cannot perform object detection.")
        return None, None # Return None for both outputs

    image = None
    input_value = None

    if input_type == "Upload File" and uploaded_image is not None:
        image = uploaded_image  # This is a PIL Image
        print("Using uploaded image (PIL) for object detection")  # Debug print

    elif input_type == "Enter URL" and image_url and image_url.strip():
        input_value = image_url
        print(f"Using URL for object detection: {input_value}")  # Debug print

    elif input_type == "Enter Base64" and base64_string and base64_string.strip():
        input_value = base64_string
        print(f"Using Base64 string for object detection")  # Debug print

    else:
        print("No valid input provided for object detection based on selected type.")
        return None, None # Return None for both outputs

    # If input_value is set (URL or Base64), use load_image
    if input_value:
        image = load_image(input_value)
        if image is None:
            return None, None # load_image failed

    # Now 'image' should be a PIL Image or None
    if image is None:
        print("Image is None after loading/selection for object detection.")
        return None, None # Return None for both outputs

    try:
        # Preprocess the image (convert PIL to numpy, ensure RGB)
        processed_image_np = preprocess_image(image)

        # Perform inference
        results = model.predict(processed_image_np)

        # Extract raw detection data
        raw_data = []
        if results and results[0].boxes:
            for box in results[0].boxes:
                # box.xywh contains [x_center, y_center, width, height]
                # box.conf contains confidence score
                # box.cls contains class index
                x_center, y_center, width, height = [round(float(coord), 2) for coord in box.xywh[0].tolist()] # Changed to xywh
                confidence = round(float(box.conf[0]), 4)
                class_id = int(box.cls[0])
                class_name = model.names[class_id] if model.names else str(class_id) # Get class name if available

                raw_data.append({
                    "box": {"x": x_center, "y": y_center, "w": width, "h": height}, # Updated keys
                    "confidence": confidence,
                    "class_id": class_id,
                    "class_name": class_name
                })

        # Draw results on the image
        result_image_np = results[0].plot() if results else processed_image_np # Plot if results exist

        print("Object detection performed successfully.")
        return result_image_np, raw_data # Return both the image and raw data

    except Exception as e:
        print(f"Error during YOLOv8 object detection: {e}")
        return None, None # Return None for both outputs