import re
from typing import Any, Callable, Sequence, TypeVar

import numpy as np
from datasets import Dataset
from smolagents import Model

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmark
from screensuite.benchmarks.perception.visualwebbench.config import (
    VisualWebBenchConfig,
    VisualWebBenchTask,
)
from screensuite.benchmarks.perception.visualwebbench.prompt import VisualWebBenchPrompt
from screensuite.benchmarks.perception.visualwebbench.utils import (
    eval_element_ocr,
    eval_element_or_action,
    eval_heading_ocr_or_web_caption,
    eval_webqa,
)
from screensuite.benchmarks.utils import bootstrap_confidence_interval
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,
)

T_RESULTS = dict[VisualWebBenchTask, float]

T_Golds = TypeVar("T_Golds", list[str], list[list[str]], list[int])
MetricFunc = Callable[[list[str], T_Golds], float]


class VisualWebBenchBenchmark(HubBaseBenchmark[VisualWebBenchConfig, None]):
    def __init__(self, name: str, config: VisualWebBenchConfig, tags: list[str]):
        """
        Initialize the VisualWebBench benchmark

        Args:
            name: Name of the benchmark.
            config: Configuration for the benchmark.
            tags: Tags for the benchmark.
        """

        super().__init__(
            name=name,
            config=config,
            tags=tags,
        )
        self.prompt_formatter: Callable[[dict], str] | None = None

        self._prompt_formatters: dict[VisualWebBenchTask, Callable[[dict], str]] = {
            VisualWebBenchTask.CAPTION_TASK: lambda _: VisualWebBenchPrompt.WEB_CAPTION_PROMPT.value,
            VisualWebBenchTask.HEADING_OCR_TASK: lambda _: VisualWebBenchPrompt.HEADING_OCR_PROMPT.value,
            VisualWebBenchTask.WEBQA_TASK: lambda sample: VisualWebBenchPrompt.WEBQA_PROMPT.value.format(
                question=sample["question"]
            ),
            VisualWebBenchTask.ELEMENT_OCR_TASK: lambda sample: VisualWebBenchPrompt.ELEMENT_OCR_PROMPT.value.format(
                bbox_ratio=sample["bbox"]
            ),
            VisualWebBenchTask.ELEMENT_GROUND_TASK: lambda sample: VisualWebBenchPrompt.ELEMENT_GROUND_PROMPT.value.format(
                element_desc=sample["elem_desc"]
            ),
            VisualWebBenchTask.ACTION_PREDICTION_TASK: lambda sample: VisualWebBenchPrompt.ACTION_PREDICTION_PROMPT.value.format(
                bbox_ratio=sample["bbox"], choices_text=sample["options"]
            ),
            VisualWebBenchTask.ACTION_GROUND_TASK: lambda sample: VisualWebBenchPrompt.ACTION_GROUND_PROMPT.value.format(
                instruction=sample["instruction"]
            ),
        }

        self._metrics: dict[VisualWebBenchTask, MetricFunc] = {
            VisualWebBenchTask.CAPTION_TASK: eval_heading_ocr_or_web_caption,
            VisualWebBenchTask.HEADING_OCR_TASK: eval_heading_ocr_or_web_caption,
            VisualWebBenchTask.WEBQA_TASK: eval_webqa,
            VisualWebBenchTask.ELEMENT_OCR_TASK: eval_element_ocr,
            VisualWebBenchTask.ELEMENT_GROUND_TASK: eval_element_or_action,
            VisualWebBenchTask.ACTION_PREDICTION_TASK: eval_element_or_action,
            VisualWebBenchTask.ACTION_GROUND_TASK: eval_element_or_action,
        }

    def set_system_prompt(self, task_type: VisualWebBenchTask) -> None:
        self.prompt_formatter = self._prompt_formatters.get(task_type)
        if self.prompt_formatter is None:
            raise NotImplementedError(f"Task type {task_type} not implemented.")

    def _get_annotated_input_from_sample(
        self, sample: dict[str, Any], evaluation_config: EvaluationConfig
    ) -> AnnotatedInput[T_Golds]:
        """
        Transforms the sample into a list of input messages and a ground truth.
        """
        assert self.prompt_formatter is not None
        prompt = self.prompt_formatter(sample)

        if evaluation_config.image_before_text:
            role_message = ChatMessage(
                role="user",
                content=[ImageContent(image=sample["image"]), TextContent(text=prompt)],
            )
        else:
            role_message = ChatMessage(
                role="user",
                content=[TextContent(text=prompt), ImageContent(image=sample["image"])],
            )

        messages = dump_chat_message_list([role_message])

        return AnnotatedInput(messages=messages, ground_truth=sample["answer"])

    def _process_response(self, response_content: str | None, task_type: VisualWebBenchTask) -> str:
        """
        Process the model response based on the task type.

        Args:
            response_content: The content of the model response
            task_type: The type of task being evaluated

        Returns:
            Processed response string
        """
        if response_content is None:
            return ""

        if task_type == VisualWebBenchTask.CAPTION_TASK:
            pattern = re.compile(r"<meta name=\"description\" content=\"(.*)\">")
            cur_meta = re.findall(pattern, response_content)
            if cur_meta:
                return cur_meta[0]
            return response_content

        elif task_type == VisualWebBenchTask.ACTION_PREDICTION_TASK:
            return response_content[0].upper() if response_content else ""

        elif task_type in [VisualWebBenchTask.WEBQA_TASK, VisualWebBenchTask.ELEMENT_OCR_TASK]:
            if ":" not in response_content:
                return response_content

            processed = ":".join(response_content.split(":")[1:])
            return processed.strip().strip('"').strip("'")

        else:
            return response_content

    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
        """

        if not self.datasets:
            self.load()

        results: T_RESULTS = {}

        for task_type in VisualWebBenchTask:
            preds: list[str] = []
            golds: list[str] | list[list[str]] | list[int] = []

            dataset = self.datasets[task_type.value]

            assert isinstance(dataset, (Dataset, Sequence))

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

            for response in responses:
                if response is not None:
                    preds.append(self._process_response(response.output, task_type))
                    golds.append(response.ground_truth)

            # Calculate metrics for this task type
            results[task_type] = self._metrics[task_type](preds, golds)

        results_dict: dict[str, float] = {task_type.value: eval_score for task_type, eval_score in results.items()}
        results_dict["average"] = float(np.mean(list(results.values())))
        results_dict["proportion_missing"] = self._calculate_proportion_missing(responses)
        results_dict["count_samples"] = len(responses)
        _, (lower, upper) = bootstrap_confidence_interval(list(results.values()))
        results_dict["average_confidence_interval_lower"] = float(lower)
        results_dict["average_confidence_interval_upper"] = float(upper)

        return BenchmarkResult(results_dict, "average")


visualwebbench = VisualWebBenchBenchmark(
    name="visualwebbench",
    config=VisualWebBenchConfig.create(),
    tags=["visualwebbench", "hf_dataset", "click", "ocr", "grounding", "caption", "webqa", "to_evaluate"],
)

# test
# from smolagents import OpenAIServerModel
#
# model = OpenAIServerModel(model_id="gpt-4o-mini")
# results = visualwebbench.evaluate(model)

registry.register(visualwebbench)
