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.perception.screenqa.config import ScreenQaConfig
from screensuite.benchmarks.perception.screenqa.utils import NO_ANSWER, sqa_s_metrics
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 AnnotatedOutput, get_model_responses

logger = logging.getLogger(__name__)


class ScreenQABenchmark(HubBaseBenchmark[ScreenQaConfig, None]):
    def _get_annotated_input_from_sample(
        self, sample: dict[str, Any], evaluation_config: EvaluationConfig
    ) -> AnnotatedInput[list[str]]:
        """
        Transform the sample into an annotated input for the model.
        """
        prompt = self.config.system_prompt.format(question=sample["question"], no_answer=NO_ANSWER)

        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["ground_truth"])

    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

        Returns:
            Evaluation results
        """

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

        exact_match_scores: list[int] = []
        f1_scores: list[float] = []

        responses: list[AnnotatedOutput[list[str]] | 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 not None and annotated_output.output is not None:
                pred, gold_list = annotated_output.output, annotated_output.ground_truth
                exact_match, f1 = sqa_s_metrics(pred, gold_list)  # type: ignore
                exact_match_scores.append(exact_match)
                f1_scores.append(f1)

        _, (lower, upper) = bootstrap_confidence_interval(f1_scores)
        results_dict: dict[str, float] = {
            "exact_match": float(np.mean(exact_match_scores)),
            "f1": float(np.mean(f1_scores)),
            "proportion_missing": self._calculate_proportion_missing(responses),
            "count_samples": len(responses),
            "f1_confidence_interval_lower": float(lower),
            "f1_confidence_interval_upper": float(upper),
        }

        return BenchmarkResult(results_dict, "f1")


screenqa_short = ScreenQABenchmark(
    name="screenqa_short",
    config=ScreenQaConfig.short(),
    tags=["screenqa", "hf_dataset", "webqa", "short", "mobile", "to_evaluate"],
)

screenqa_complex = ScreenQABenchmark(
    name="screenqa_complex",
    config=ScreenQaConfig.complex(),
    tags=["screenqa", "hf_dataset", "webqa", "complex", "mobile", "to_evaluate"],
)

registry.register([screenqa_short, screenqa_complex])


# if __name__ == "__main__":
#     # test
#     from smolagents import OpenAIServerModel
#
#     model = OpenAIServerModel(model_id="gpt-4o-mini")
#     screenqa_short.evaluate(model, EvaluationConfig(timeout=10, parallel_workers=32, test_mode=True))
#     screenqa_complex.evaluate(model, EvaluationConfig(timeout=10, parallel_workers=32, test_mode=True))
