import unittest
from unittest.mock import patch

from PIL import Image
from smolagents import ChatMessage

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmarks.model import BoundingBox
from screensuite.benchmarks.perception.screenspot.benchmark import ScreenSpotBenchmark
from screensuite.benchmarks.perception.screenspot.config import (
    ScreenSpotConfig,
    ScreenSpotDatasetVersion,
    ScreenSpotHFRepo,
)
from screensuite.benchmarks.perception.screenspot.prompt import LocalizationPrompt


class TestScreenSpotBenchmark(unittest.TestCase):
    """Test cases for the ScreenSpotBenchmark class."""

    def setUp(self):
        """Set up test fixtures."""
        # Create a mock config
        self.config = ScreenSpotConfig(
            dataset_version=ScreenSpotDatasetVersion.V1,
            hf_repo=ScreenSpotHFRepo.V1,
            split="test",
            system_prompt=LocalizationPrompt.CLICK_PROMPT_NORMALIZED,
            accuracy_threshold=0.8,
        )

        # Create a benchmark instance
        self.benchmark = ScreenSpotBenchmark(
            name="test-benchmark",
            config=self.config,
            tags=["test"],
        )

        # Create a mock dataset
        self.mock_dataset = [
            {
                "image": Image.new("RGB", (100, 100), color="white"),
                "instruction": "Click on the button",
                "bbox": [0.1, 0.1, 0.3, 0.3],
            },
            {
                "image": Image.new("RGB", (100, 100), color="white"),
                "instruction": "Click on the link",
                "bbox": [0.5, 0.5, 0.7, 0.7],
            },
        ]

    @patch("screensuite.benchmarks.hubbasebenchmark.login")
    @patch("screensuite.benchmarks.hubbasebenchmark.load_dataset")
    def test_load(self, mock_load_dataset, mock_login):
        """Test the load method."""
        # Set up the mock
        mock_load_dataset.return_value = self.mock_dataset

        # Call the method
        self.benchmark.load()

        # Check that the result is the mock dataset
        self.assertEqual(self.benchmark.dataset, self.mock_dataset)

        # Check that the dataset was stored in the instance
        self.assertEqual(self.benchmark.dataset, self.mock_dataset)

        # Test that the method returns the cached dataset on subsequent calls
        mock_load_dataset.reset_mock()
        self.benchmark.load()
        mock_load_dataset.assert_not_called()
        self.assertEqual(self.benchmark.dataset, self.mock_dataset)

    def test_calculate_accuracy_click_inside(self):
        """Test the _calculate_accuracy method for click prompts with a point inside the box."""

        self.benchmark.config.system_prompt = LocalizationPrompt.CLICK_PROMPT_NORMALIZED

        # Test with a point inside the box
        predicted = "Click(0.2, 0.2)"
        result = self.benchmark._calculate_accuracy(
            predicted, (BoundingBox.from_xyxy(xyxy_bbox=(0.1, 0.1, 0.3, 0.3)), (1, 1))
        )
        self.assertEqual(result, 1.0)

    def test_calculate_accuracy_click_outside(self):
        """Test the _calculate_accuracy method for click prompts with a point outside the box."""

        self.benchmark.config.system_prompt = LocalizationPrompt.CLICK_PROMPT_NORMALIZED

        # Test with a point outside the box
        predicted = "Click(0.5, 0.5)"
        result = self.benchmark._calculate_accuracy(
            predicted, (BoundingBox.from_xyxy(xyxy_bbox=(0.1, 0.1, 0.3, 0.3)), (1, 1))
        )
        self.assertEqual(result, 0.0)

    def test_calculate_accuracy_bounding_box_overlap(self):
        """Test the _calculate_accuracy method for bounding box prompts with overlapping boxes."""

        self.benchmark.config.system_prompt = LocalizationPrompt.BOUNDING_BOX_PROMPT

        # Test with overlapping boxes
        predicted = "BoundingBox(0.2, 0.2, 0.4, 0.4)"
        result = self.benchmark._calculate_accuracy(
            predicted, (BoundingBox.from_xyxy(xyxy_bbox=(0.1, 0.1, 0.3, 0.3)), (1, 1))
        )
        self.assertGreater(result, 0.0)

    def test_calculate_accuracy_bounding_box_no_overlap(self):
        """Test the _calculate_accuracy method for bounding box prompts with non-overlapping boxes."""

        self.benchmark.config.system_prompt = LocalizationPrompt.BOUNDING_BOX_PROMPT

        # Test with non-overlapping boxes
        predicted = "BoundingBox(0.5, 0.5, 0.7, 0.7)"
        result = self.benchmark._calculate_accuracy(
            predicted, (BoundingBox.from_xyxy(xyxy_bbox=(0.1, 0.1, 0.3, 0.3)), (1, 1))
        )
        self.assertEqual(result, 0.0)

    def test_calculate_accuracy_bounding_box_perfect_overlap(self):
        """Test the _calculate_accuracy method for bounding box prompts with perfectly overlapping boxes."""

        self.benchmark.config.system_prompt = LocalizationPrompt.BOUNDING_BOX_PROMPT

        # Test with perfectly overlapping boxes
        predicted = "BoundingBox(0.1, 0.1, 0.3, 0.3)"
        result = self.benchmark._calculate_accuracy(
            predicted, (BoundingBox.from_xyxy(xyxy_bbox=(0.1, 0.1, 0.3, 0.3)), (1, 1))
        )
        self.assertEqual(result, 1.0)

    @patch("screensuite.benchmarks.perception.screenspot.benchmark.Model")
    @patch("screensuite.benchmarks.hubbasebenchmark.load_dataset")
    def test_evaluate(self, mock_load_dataset, mock_model):
        """Test the evaluate method."""
        # Set up the mock model
        self.benchmark.config.system_prompt = LocalizationPrompt.CLICK_PROMPT_NORMALIZED
        mock_model.generate.return_value = ChatMessage(role="assistant", content="click(0.5, 0.5)")
        mock_load_dataset.return_value = self.mock_dataset
        # Call the method
        results = self.benchmark.evaluate(mock_model, EvaluationConfig(timeout=10, parallel_workers=1, run_name=None))
        self.assertEqual(
            results._metrics,
            {
                "avg_accuracy": 0.5,
                "avg_accuracy_confidence_interval_lower": 0.0,
                "avg_accuracy_confidence_interval_upper": 1.0,
                "success_rate": 0.5,
                "proportion_missing": 0.0,
                "count_samples": 2,
            },
        )

        self.benchmark.config.system_prompt = LocalizationPrompt.BOUNDING_BOX_PROMPT
        mock_model.generate.return_value = ChatMessage(role="assistant", content="BoundingBox(0.1, 0.1, 0.3, 0.3)")
        results = self.benchmark.evaluate(mock_model, EvaluationConfig(timeout=10, parallel_workers=1, run_name=None))
        self.assertEqual(
            results._metrics,
            {
                "avg_accuracy": 0.5,
                "avg_accuracy_confidence_interval_lower": 0.0,
                "avg_accuracy_confidence_interval_upper": 1.0,
                "success_rate": 0.5,
                "proportion_missing": 0.0,
                "count_samples": 2,
            },
        )


if __name__ == "__main__":
    unittest.main()
