import os
import time
from abc import ABC, abstractmethod
from typing import Any

import numpy as np
from PIL import Image, ImageDraw
from smolagents import CodeAgent, Model, Tool
from smolagents.agent_types import AgentImage
from smolagents.memory import ActionStep, TaskStep
from smolagents.monitoring import LogLevel


class VisionAgent(CodeAgent, ABC):
    """Base class for vision-based agents that interact with visual environments."""

    def __init__(
        self,
        model: Model,
        data_dir: str,
        tools: list[Tool] | None = None,
        max_steps: int = 200,
        verbosity_level: LogLevel = LogLevel.INFO,
        planning_interval: int | None = None,
        **kwargs,
    ):
        """Initialize the vision agent.

        Args:
            model: The model to use for vision and reasoning
            data_dir: Directory to save screenshots and other data
            tools: Optional list of tools to use
            max_steps: Maximum number of steps to run
            verbosity_level: Level of logging verbosity
            planning_interval: Interval between planning steps
            **kwargs: Additional arguments passed to parent class
        """
        self._data_dir = data_dir
        self.planning_interval = planning_interval

        # Set up temp directory
        os.makedirs(self._data_dir, exist_ok=True)
        print(f"Screenshots and steps will be saved to: {self._data_dir}")

        # Initialize base agent
        super().__init__(
            tools=tools or [],
            model=model,
            max_steps=max_steps,
            verbosity_level=verbosity_level,
            planning_interval=self.planning_interval,
            **kwargs,
        )

        # Add default tools
        self.logger.log("Setting up agent tools...")
        self._setup_desktop_tools()
        self.step_callbacks.append(self.take_screenshot_callback)
        self.click_coordinates: list[int] | None = None

    @abstractmethod
    def _setup_desktop_tools(self) -> None:
        """Set up the tools specific to this agent's environment."""
        ...

    @abstractmethod
    def get_screenshot(self) -> np.ndarray[Any, Any]:
        """Get a screenshot from the environment.

        Returns:
            numpy array containing the screenshot
        """
        ...

    @property
    def data_dir(self) -> str:
        return self._data_dir

    @data_dir.setter
    def data_dir(self, data_dir: str) -> None:
        self._data_dir = data_dir
        os.makedirs(self._data_dir, exist_ok=True)

    def draw_marker_on_image(self, image: Image.Image, click_coordinates: list[int] | None) -> Image.Image:
        """Draw a marker on an image at the specified coordinates.

        Args:
            image: The image to draw on
            click_coordinates: The coordinates where the marker should be drawn

        Returns:
            The image with the marker drawn on it
        """
        x, y = click_coordinates or (0, 0)
        draw = ImageDraw.Draw(image)
        cross_size, linewidth = 10, 3
        # Draw cross
        draw.line((x - cross_size, y, x + cross_size, y), fill="green", width=linewidth)
        draw.line((x, y - cross_size, x, y + cross_size), fill="green", width=linewidth)
        # Add a circle around it for better visibility
        draw.ellipse(
            (
                x - cross_size * 2,
                y - cross_size * 2,
                x + cross_size * 2,
                y + cross_size * 2,
            ),
            outline="green",
            width=linewidth,
        )
        return image

    def take_screenshot_callback(self, memory_step: ActionStep, agent: CodeAgent) -> None:
        """Callback that takes a screenshot + memory snapshot after a step completes.

        Args:
            memory_step: The current memory step
            agent: The agent instance
        """
        self.logger.log("Analyzing screen content...")

        current_step = memory_step.step_number or 0

        time.sleep(2.5)  # Let things happen on the desktop
        # Get screenshot from environment
        screenshot_array = self.get_screenshot()
        image = Image.fromarray(screenshot_array)

        # Create a filename with step number
        screenshot_path = os.path.join(self._data_dir, f"step_{current_step:03d}.png")
        image.save(screenshot_path)

        image_copy = image.copy()

        if self.click_coordinates is not None:
            image_copy = self.draw_marker_on_image(image_copy, self.click_coordinates)

        self.last_marked_screenshot = AgentImage(screenshot_path)
        print(f"Saved screenshot for step {current_step} to {screenshot_path}")

        # Clean up previous screenshots from memory
        for previous_memory_step in agent.memory.steps:
            if (
                isinstance(previous_memory_step, ActionStep)
                and (previous_memory_step.step_number or 0) <= current_step - 1
            ):
                previous_memory_step.observations_images = None
            elif isinstance(previous_memory_step, TaskStep):
                previous_memory_step.task_images = None

            # Check for repeated actions
            if (
                isinstance(previous_memory_step, ActionStep)
                and (previous_memory_step.step_number or 0) == current_step - 1
            ):
                if (
                    previous_memory_step.tool_calls
                    and getattr(previous_memory_step.tool_calls[0], "arguments", None)
                    and memory_step.tool_calls
                    and getattr(memory_step.tool_calls[0], "arguments", None)
                ):
                    if previous_memory_step.tool_calls[0].arguments == memory_step.tool_calls[0].arguments:
                        if isinstance(memory_step.observations, str):
                            memory_step.observations += "\nWARNING: You've executed the same action several times in a row. MAKE SURE TO NOT UNNECESSARILY REPEAT ACTIONS."
                        else:
                            memory_step.observations = "WARNING: You've executed the same action several times in a row. MAKE SURE TO NOT UNNECESSARILY REPEAT ACTIONS."

        # Add the screenshot to the current memory step
        memory_step.observations_images = [image_copy]
        # memory_step.observations_images = [screenshot_path] # IF YOU USE THIS INSTEAD OF ABOVE, LAUNCHING A SECOND TASK BREAKS

        self.click_coordinates = None  # Reset click marker

    @abstractmethod
    def close(self) -> None:
        """Clean up resources used by the agent."""
        pass
