import base64
import logging
from io import BytesIO
from typing import Any

import numpy as np
from PIL import Image
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.websrc.config import WebSrcConfig
from screensuite.benchmarks.utils import (
    bootstrap_confidence_interval,
    compute_exact,
    compute_f1,
)
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 WebSrcBenchmark(HubBaseBenchmark[WebSrcConfig, None]):
    def _get_annotated_input_from_sample(
        self, sample: dict[str, Any], evaluation_config: EvaluationConfig
    ) -> AnnotatedInput[str]:
        """
        Transforms the sample into a list of input messages and a ground truth.
        """
        prompt = self.config.system_prompt.format(question=sample["question"])
        # Create validated PIL image from the base64 string
        image = Image.open(BytesIO(base64.b64decode(sample["image"])))

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

        messages = dump_chat_message_list([role_message])

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

    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[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,
        )
        print("HERE ARE THE RESPONSES", responses)

        for response in responses:
            if response is not None:
                pred, gold = response.output, response.ground_truth
                exact_match, f1 = compute_exact(pred, gold), compute_f1(pred, gold)
                exact_match_scores.append(exact_match)
                f1_scores.append(f1)

        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),
        }
        _, (lower, upper) = bootstrap_confidence_interval(f1_scores)
        results_dict["f1_confidence_interval_lower"] = float(lower)
        results_dict["f1_confidence_interval_upper"] = float(upper)

        return BenchmarkResult(results_dict, "f1")


websrc_dev = WebSrcBenchmark(
    name="websrc_dev",
    config=WebSrcConfig.dev(),
    tags=["websrc", "hf_dataset", "webqa", "dev", "to_evaluate"],
)


registry.register(websrc_dev)


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