#!/usr/bin/env python3
"""Object detection module using Google Gemini API"""

import json
import logging
import os
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Any
from PIL import Image, ImageDraw

from google import genai
from google.genai import types


logger = logging.getLogger(__name__)


def detect_objects(
    image_path: str,
    prompt: str,
    api_key: Optional[str] = None,
    model: str = "gemini-2.0-flash-exp"
) -> List[Dict[str, Any]]:
    """
    Detect objects in an image using Google Gemini API.

    Args:
        image_path: Path to the input image
        prompt: Prompt for object detection (required)
        api_key: Google Gemini API key (optional, can use env var)
        model: Model to use for detection

    Returns:
        List of detected objects with bounding boxes
    """
    # Prefer GEMINI_API_KEY from environment; no additional fallback or validation.
    if api_key is None:
        api_key = os.environ.get('GEMINI_API_KEY')

    # Check if image exists
    image_path = Path(image_path)
    if not image_path.exists():
        raise FileNotFoundError(f"Image not found: {image_path}")

    # Initialize client using env-provided key if present
    client = genai.Client(api_key=api_key) if api_key else genai.Client()

    # Ensure prompt includes bounding box format requirement
    if "box_2d" not in prompt.lower() and "bounding" not in prompt.lower():
        prompt += " Return the result as JSON array where each object has 'name' and 'box_2d' fields. The box_2d should be [ymin, xmin, ymax, xmax] normalized to 0-1000."

    # Open image
    image = Image.open(image_path)

    # Configure JSON response
    config = types.GenerateContentConfig(
        response_mime_type="application/json"
    )

    logger.info(f"Detecting objects in {image_path} using {model}")

    # Generate response
    try:
        response = client.models.generate_content(
            model=model,
            contents=[image, prompt],
            config=config
        )

        # Parse response
        bounding_boxes = json.loads(response.text)
        logger.debug(f"Raw response: {bounding_boxes}")

        # Convert normalized coordinates to absolute
        width, height = image.size
        converted_objects = []

        for obj in bounding_boxes:
            if isinstance(obj, dict) and "box_2d" in obj:
                box = obj["box_2d"]
                # Ensure box has 4 coordinates
                if not isinstance(box, (list, tuple)) or len(box) < 4:
                    logger.warning(f"Invalid box format: {box}")
                    continue

                abs_y1 = int(box[0] / 1000 * height)
                abs_x1 = int(box[1] / 1000 * width)
                abs_y2 = int(box[2] / 1000 * height)
                abs_x2 = int(box[3] / 1000 * width)

                converted_obj = {
                    "name": obj.get("name", "unknown"),
                    "box": [abs_x1, abs_y1, abs_x2, abs_y2],
                    "normalized_box": box,
                    "confidence": obj.get("confidence", None)
                }
                converted_objects.append(converted_obj)

        logger.info(f"Detected {len(converted_objects)} objects")
        return converted_objects

    except Exception as e:
        logger.error(f"Error detecting objects: {e}")
        raise


def draw_bounding_boxes(
    image_path: str,
    objects: List[Dict[str, Any]],
    output_path: Optional[str] = None,
    box_color: str = "red",
    box_width: int = 1,
    dash_length: int = 5
) -> str:
    """
    Draw bounding boxes on an image.

    Args:
        image_path: Path to the input image
        objects: List of detected objects with bounding boxes
        output_path: Path to save the output image
        box_color: Color of the bounding boxes
        box_width: Width of the bounding box lines
        dash_length: Length of dashes for dashed lines

    Returns:
        Path to the output image
    """
    # Open image
    image = Image.open(image_path)
    draw = ImageDraw.Draw(image)

    # Draw bounding boxes
    for obj in objects:
        if "box" in obj:
            x1, y1, x2, y2 = obj["box"]

            # Draw dashed rectangle
            # Top edge
            for i in range(x1, x2, dash_length * 2):
                draw.line(
                    [(i, y1), (min(i + dash_length, x2), y1)],
                    fill=box_color,
                    width=box_width
                )

            # Bottom edge
            for i in range(x1, x2, dash_length * 2):
                draw.line(
                    [(i, y2), (min(i + dash_length, x2), y2)],
                    fill=box_color,
                    width=box_width
                )

            # Left edge
            for i in range(y1, y2, dash_length * 2):
                draw.line(
                    [(x1, i), (x1, min(i + dash_length, y2))],
                    fill=box_color,
                    width=box_width
                )

            # Right edge
            for i in range(y1, y2, dash_length * 2):
                draw.line(
                    [(x2, i), (x2, min(i + dash_length, y2))],
                    fill=box_color,
                    width=box_width
                )

            # Add label if available
            if "name" in obj and obj["name"] != "unknown":
                # Draw text background
                text = obj["name"]
                # Note: PIL doesn't have built-in text size calculation without font
                # For simplicity, we'll just draw the text without background
                draw.text((x1, y1 - 15), text, fill=box_color)

    # Save output
    if not output_path:
        # Generate output path based on input
        input_path = Path(image_path)
        output_path = input_path.parent / f"{input_path.stem}_detected{input_path.suffix}"

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    image.save(output_path)

    logger.info(f"Saved image with bounding boxes to {output_path}")
    return str(output_path)


def detect_and_draw(
    image_path: str,
    prompt: str,
    output_path: Optional[str] = None,
    api_key: Optional[str] = None,
    model: str = "gemini-2.0-flash-exp",
    box_color: str = "red",
    box_width: int = 1
) -> Tuple[str, List[Dict[str, Any]]]:
    """
    Detect objects and draw bounding boxes on the image.

    Args:
        image_path: Path to the input image
        prompt: Prompt for object detection (required)
        output_path: Path to save the output image
        api_key: Google Gemini API key
        model: Model to use for detection
        box_color: Color of the bounding boxes
        box_width: Width of the bounding box lines

    Returns:
        Tuple of (output image path, detected objects)
    """
    # Detect objects
    objects = detect_objects(image_path, prompt, api_key, model)

    # Draw bounding boxes
    output_file = draw_bounding_boxes(
        image_path,
        objects,
        output_path,
        box_color=box_color,
        box_width=box_width
    )

    return output_file, objects
