import logging
import os
from typing import Any, Generator

import numpy as np
from datasets import load_dataset
from PIL import Image
from smolagents import Model

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmark
from screensuite.benchmarks.model import BoundingBox, NormalizedClickCoordinates
from screensuite.benchmarks.singlestep.showdown_clicks.config import (
    ShowdownClicksConfig,
)
from screensuite.benchmarks.singlestep.showdown_clicks.prompt import (
    ShowdownClicksPrompt,
)
from screensuite.benchmarks.utils import (
    bootstrap_confidence_interval,
    evaluate_click_distance,
    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 (
    AnnotatedInput,
    AnnotatedOutput,
    get_model_responses,
)

logger = logging.getLogger(__name__)

ActionGroundTruth = tuple[BoundingBox, tuple[int, int]]


class ShowdownClicksBenchmark(HubBaseBenchmark[ShowdownClicksConfig, None]):
    def load(self, streaming: bool = False) -> None:
        from huggingface_hub import snapshot_download

        snapshot_download(
            repo_id="generalagents/showdown-clicks",
            repo_type="dataset",
            local_dir="downloaded_data",
            ignore_patterns=[".gitattributes", "README.md"],
        )
        self.dataset = load_dataset(self.config.hf_repo, split=self.config.split, streaming=streaming)

        # # Add an "image" column to the dataset with PIL Images loaded from the frame paths

        def load_image(sample):
            image_path = os.path.join("downloaded_data/showdown-clicks-dev", sample["image"])

            image = Image.open(image_path)
            return {"screenshot": image}

        self.dataset = self.dataset.map(load_image)

    def _get_annotated_input_from_sample(
        self, sample_dict: dict[str, list[Any] | str], evaluation_config: EvaluationConfig
    ) -> Generator[AnnotatedInput[ActionGroundTruth], None, None]:
        """
        Transform the sample_dict into an AnnotatedInput with a list of messages and a ground truth.
        """
        target_bounding_box = BoundingBox.from_xyxy(
            xyxy_bbox=(
                float(sample_dict["x1"]),  # type: ignore
                float(sample_dict["y1"]),  # type: ignore
                float(sample_dict["x2"]),  # type: ignore
                float(sample_dict["y2"]),  # type: ignore
            ),
            image_width=int(sample_dict["width"]),  # type: ignore
            image_height=int(sample_dict["height"]),  # type: ignore
        )
        image_dimensions = (int(sample_dict["width"]), int(sample_dict["height"]))  # type: ignore

        if evaluation_config.use_uitars_action_space:
            prompt = ShowdownClicksPrompt.PROMPT_ABSOLUTE_CLICK_UITARS.value.format(
                instruction=sample_dict["instruction"], resolution=image_dimensions
            )
        elif evaluation_config.use_h_action_space:
            prompt = ShowdownClicksPrompt.PROMPT_ABSOLUTE_CLICK_H.value.format(instruction=sample_dict["instruction"])
        else:
            prompt = ShowdownClicksPrompt.PROMPT_ABSOLUTE_CLICK.value.format(
                instruction=sample_dict["instruction"], resolution=image_dimensions
            )

        optional_system_prompt = (
            [ChatMessage(role="system", content=[TextContent(text=evaluation_config.custom_system_prompt)])]
            if evaluation_config.custom_system_prompt
            else []
        )

        if evaluation_config.image_before_text:
            input_message = ChatMessage(
                role="user",
                content=[
                    ImageContent(
                        image=sample_dict["screenshot"],  # type: ignore
                    ),
                    TextContent(
                        text=prompt,
                    ),
                ],
            )
        else:
            input_message = ChatMessage(
                role="user",
                content=[
                    TextContent(
                        type="text",
                        text=prompt,
                    ),
                    ImageContent(
                        image=sample_dict["screenshot"],  # type: ignore
                    ),
                ],
            )
        yield AnnotatedInput(
            messages=dump_chat_message_list(optional_system_prompt + [input_message]),
            ground_truth=(target_bounding_box, image_dimensions),
        )

    def score_single_output(
        self,
        annotated_output: AnnotatedOutput[ActionGroundTruth],
    ) -> tuple[float, float]:
        prediction, (target_bounding_box, image_dimensions) = (
            annotated_output.output,
            annotated_output.ground_truth,
        )

        target_click_coordinates = NormalizedClickCoordinates.from_xy(
            xy=(target_bounding_box.x, target_bounding_box.y),
            image_width=image_dimensions[0],
            image_height=image_dimensions[1],
        )
        click_acc = evaluate_click_distance(
            prediction, target_click_coordinates, image_width=image_dimensions[0], image_height=image_dimensions[1]
        )
        in_bounding_box = evaluate_click_in_bounding_box(
            prediction, target_bounding_box, image_width=image_dimensions[0], image_height=image_dimensions[1]
        )
        return click_acc, in_bounding_box

    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: Configuration for inference

        Returns:
            Evaluation results
        """
        if not self.datasets:
            self.load(streaming=False)

        accuracy_scores: list[float] = []
        in_bounding_box_scores: list[float] = []
        responses: list[AnnotatedOutput[ActionGroundTruth] | None] = []

        # self._aggregate_traces()
        responses = get_model_responses(
            self.dataset,  # type: ignore
            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 not None:
                score, in_bounding_box = self.score_single_output(annotated_output)
                accuracy_scores.append(score)
                in_bounding_box_scores.append(in_bounding_box)
        metrics = {
            "bounding_box_acc": float(np.mean(in_bounding_box_scores)) if in_bounding_box_scores else np.nan,
            "click_acc": float(np.mean(accuracy_scores)) if accuracy_scores else np.nan,
        }
        reference_field = "bounding_box_acc"
        _, (lower, upper) = bootstrap_confidence_interval(accuracy_scores)
        metrics["action_acc_confidence_interval_lower"] = float(lower)
        metrics["action_acc_confidence_interval_upper"] = float(upper)
        metrics["proportion_missing"] = self._calculate_proportion_missing(responses)
        metrics["count_samples"] = len(responses)
        return BenchmarkResult(metrics=metrics, reference_field=reference_field)


showdown_clicks = ShowdownClicksBenchmark(
    name="showdown_clicks",
    config=ShowdownClicksConfig(),
    tags=["showdown_clicks", "hf_dataset", "offline", "agent", "web", "trace", "to_evaluate"],
)


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

    model = OpenAIServerModel(model_id="gpt-4o")
    print(
        showdown_clicks.evaluate(
            model, EvaluationConfig(timeout=10, parallel_workers=4, test_mode=True, run_name="test")
        )
    )
else:
    registry.register(showdown_clicks)
