import logging
from typing import Any

import numpy as np
from smolagents import Model

from screensuite.basebenchmark import AnnotatedInput, EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmark
from screensuite.benchmarks.model import BoundingBox
from screensuite.benchmarks.perception.screenspot.config import ScreenSpotConfig
from screensuite.benchmarks.perception.screenspot.prompt import LocalizationPrompt
from screensuite.benchmarks.utils import (
    bootstrap_confidence_interval,
    evaluate_bounding_box_action,
    evaluate_click_in_bounding_box,
)
from screensuite.chat_message import (
    ChatMessage,
    ImageContent,
    TextContent,
    dump_chat_message_list,
)
from screensuite.registry import registry
from screensuite.response_generation import AnnotatedOutput, get_model_responses

logger = logging.getLogger(__name__)


class ScreenSpotBenchmark(HubBaseBenchmark[ScreenSpotConfig, None]):
    def _get_annotated_input_from_sample(
        self, sample: dict[str, Any], evaluation_config: EvaluationConfig
    ) -> AnnotatedInput[tuple[BoundingBox, tuple[int, int]]]:
        """
        Transforms the sample into a list of input messages and a ground truth.
        """
        image_dimensions = (sample["image"].width, sample["image"].height)  # type: ignore

        if evaluation_config.use_uitars_action_space:
            prompt = LocalizationPrompt.CLICK_PROMPT_ABSOLUTE_UITARS.value.format(
                instruction=sample["instruction"], resolution=image_dimensions
            )
        elif evaluation_config.use_h_action_space:
            prompt = LocalizationPrompt.CLICK_PROMPT_ABSOLUTE.value.replace("click(", "Click(").format(
                instruction=sample["instruction"], resolution=image_dimensions
            )
        else:
            prompt = LocalizationPrompt.CLICK_PROMPT_ABSOLUTE.value.format(
                instruction=sample["instruction"], resolution=image_dimensions
            )
        if evaluation_config.image_before_text:
            input_message = ChatMessage(
                role="user",
                content=[ImageContent(image=sample["image"]), TextContent(text=prompt)],
            )
        else:
            input_message = ChatMessage(
                role="user",
                content=[TextContent(text=prompt), ImageContent(image=sample["image"])],
            )

        optional_system_prompt = (
            [ChatMessage(role="system", content=[TextContent(text=evaluation_config.custom_system_prompt)])]
            if evaluation_config.custom_system_prompt
            else []
        )
        return AnnotatedInput(
            messages=dump_chat_message_list(optional_system_prompt + [input_message]),
            ground_truth=(BoundingBox.from_xyxy(sample["bbox"]), image_dimensions),
        )

    def _calculate_accuracy(self, response: str, ground_truth: tuple[BoundingBox, tuple[int, int]]) -> float:
        """
        Calculate accuracy between predicted and ground truth coordinates

        Args:
            response: The model's response text
            ground_truth: The ground truth, a tuple of the bounding box and the image dimensions
        """
        if self.config.system_prompt == LocalizationPrompt.CLICK_PROMPT_NORMALIZED:
            return evaluate_click_in_bounding_box(response, ground_truth[0], image_width=1, image_height=1)
        elif self.config.system_prompt == LocalizationPrompt.CLICK_PROMPT_ABSOLUTE:
            image_dimensions = ground_truth[1]
            return evaluate_click_in_bounding_box(
                response, ground_truth[0], image_width=image_dimensions[0], image_height=image_dimensions[1]
            )
        elif self.config.system_prompt == LocalizationPrompt.BOUNDING_BOX_PROMPT:
            return evaluate_bounding_box_action(response, ground_truth[0])
        else:
            raise ValueError(f"Unsupported system prompt: {self.config.system_prompt}")

    def _record_result(self, accuracy: float, is_above_threshold: list[bool], accuracies: list[float]) -> None:
        """
        Record the result of a prediction

        Args:
            accuracy: The accuracy score (0.0 to 1.0)
            correct: List to append to if correct
            accuracies: List to append accuracy to
        """
        accuracies.append(accuracy)

        if accuracy >= self.config.accuracy_threshold:
            is_above_threshold.append(True)
        else:
            is_above_threshold.append(False)

    def evaluate(
        self,
        model: Model,
        evaluation_config: EvaluationConfig,
        env_config: None = None,
    ) -> BenchmarkResult:
        """
        Evaluate the model on the benchmark

        Args:
            model: The model to evaluate
            evaluation_config: The inference configuration

        Returns:
            Evaluation results
        """
        num_action: int = 0
        is_above_threshold: list[bool] = []
        accuracies: list[float] = []

        if self.dataset is None:
            self.load()

        responses: list[AnnotatedOutput[tuple[BoundingBox, tuple[int, int]]] | None] = get_model_responses(
            self.dataset,
            model,
            self._get_annotated_input_from_sample,
            self.name,
            evaluation_config,
            max_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
        )

        for annotated_output in responses:
            if annotated_output is None:
                continue
            response, (bbox, resolution) = annotated_output.output, annotated_output.ground_truth
            num_action += 1
            accuracy = self._calculate_accuracy(response, (bbox, resolution))
            self._record_result(accuracy, is_above_threshold, accuracies)

        avg_accuracy = float(np.mean(accuracies)) if accuracies else 0.0
        success_rate = sum(1 for x in is_above_threshold if x) / num_action if num_action > 0 else 0.0
        _, (lower, upper) = bootstrap_confidence_interval(accuracies)
        results = {
            "avg_accuracy": avg_accuracy,
            "success_rate": success_rate,
            "proportion_missing": self._calculate_proportion_missing(responses),
            "count_samples": len(responses),
            "avg_accuracy_confidence_interval_lower": float(lower),
            "avg_accuracy_confidence_interval_upper": float(upper),
        }

        return BenchmarkResult(results, "avg_accuracy")


screenspot_v1_click_prompt = ScreenSpotBenchmark(
    name="screenspot-v1-click-prompt",
    config=ScreenSpotConfig.v1(LocalizationPrompt.CLICK_PROMPT_ABSOLUTE),
    tags=["screenspot", "grounding", "hf_dataset", "v1", "click"],
)

screenspot_v1_bounding_box_prompt = ScreenSpotBenchmark(
    name="screenspot-v1-bounding-box-prompt",
    config=ScreenSpotConfig.v1(LocalizationPrompt.BOUNDING_BOX_PROMPT),
    tags=["screenspot", "grounding", "hf_dataset", "v1", "bounding_box"],
)

screenspot_v2_click_prompt = ScreenSpotBenchmark(
    name="screenspot-v2-click-prompt",
    config=ScreenSpotConfig.v2(LocalizationPrompt.CLICK_PROMPT_ABSOLUTE),
    tags=["screenspot", "grounding", "hf_dataset", "v2", "click", "to_evaluate"],
)

screenspot_v2_bounding_box_prompt = ScreenSpotBenchmark(
    name="screenspot-v2-bounding-box-prompt",
    config=ScreenSpotConfig.v2(LocalizationPrompt.BOUNDING_BOX_PROMPT),
    tags=["screenspot", "grounding", "hf_dataset", "v2", "bounding_box"],
)

screenspot_pro_click_prompt = ScreenSpotBenchmark(
    name="screenspot-pro-click-prompt",
    config=ScreenSpotConfig.pro(LocalizationPrompt.CLICK_PROMPT_ABSOLUTE),
    tags=["screenspot", "grounding", "hf_dataset", "pro", "click", "to_evaluate"],
)

screenspot_pro_bounding_box_prompt = ScreenSpotBenchmark(
    name="screenspot-pro-bounding-box-prompt",
    config=ScreenSpotConfig.pro(LocalizationPrompt.BOUNDING_BOX_PROMPT),
    tags=["screenspot", "grounding", "hf_dataset", "pro", "bounding_box"],
)

if __name__ == "__main__":
    from smolagents import OpenAIServerModel

    model = OpenAIServerModel(model_id="gpt-4o")
    evaluation_config = EvaluationConfig(parallel_workers=1, run_name=None)
    benchmark_result = screenspot_v1_click_prompt.evaluate(model, evaluation_config)
else:
    registry.register(
        [
            screenspot_v1_click_prompt,
            screenspot_v1_bounding_box_prompt,
            screenspot_v2_click_prompt,
            screenspot_v2_bounding_box_prompt,
            screenspot_pro_click_prompt,
            screenspot_pro_bounding_box_prompt,
        ]
    )
