import datetime
import logging
from urllib.parse import parse_qs, unquote, urlparse

import numpy as np
from smolagents import Model

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmark
from screensuite.benchmarks.multistep.config import AgentRunResult
from screensuite.benchmarks.multistep.generation import get_agent_responses
from screensuite.benchmarks.multistep.mind2web.config import Mind2WebConfig
from screensuite.registry import registry

logger = logging.getLogger()


class MatchFunction:
    """Taken from https://github.com/iMeanAI/WebCanvas/blob/b9f289128614cd99b97abd0bb9bfc3a45f0847e0/evaluate/step_score_js.py#L193"""

    def __init__(self):
        pass

    @staticmethod
    def include_match(input_answer, reference_answer) -> int:
        return 1 if reference_answer in input_answer else 0

    @staticmethod
    def exact_match(input_answer, reference_answer) -> int:
        return 1 if input_answer == reference_answer else 0


class Mind2WebBenchmark(HubBaseBenchmark[Mind2WebConfig, None]):
    @staticmethod
    def url_exact_match(input_url, reference_answer, key=False):
        """Taken from https://github.com/iMeanAI/WebCanvas/blob/b9f289128614cd99b97abd0bb9bfc3a45f0847e0/evaluate/step_score_js.py#L36"""
        if key:
            try:
                parsed_url = urlparse(input_url)
                url_params = parse_qs(parsed_url.query)
                input_answer = url_params[key][0]
            except Exception:
                return 0
        else:
            input_answer = input_url
        input_answer = unquote(input_answer)
        result_score = MatchFunction.exact_match(input_answer, reference_answer)
        return result_score

    @staticmethod
    def url_include_match(input_url, reference_answer, key=None):
        """Taken from https://github.com/iMeanAI/WebCanvas/blob/b9f289128614cd99b97abd0bb9bfc3a45f0847e0/evaluate/step_score_js.py#L36"""
        if key:
            try:
                parsed_url = urlparse(input_url)
                url_params = parse_qs(parsed_url.query)
                input_answer = url_params[key][0]
            except Exception:
                return 0
        else:
            try:
                parsed_url = urlparse(input_url)
                input_answer = parsed_url.netloc + parsed_url.path
                if parsed_url.fragment is not None and (parsed_url.fragment):
                    input_answer += "#" + parsed_url.fragment
            except Exception:
                input_answer = input_url
        input_answer = unquote(input_answer)
        result_score = MatchFunction.include_match(input_answer, reference_answer)
        return result_score

    def load(self, streaming: bool = False) -> None:
        """
        Load the dataset from Hugging Face

        Returns:
            The loaded dataset
        """
        super().load(streaming=streaming)
        assert self.dataset is not None  # type: ignore
        self.dataset = self.dataset.rename_columns({"task": "question", "evaluation": "reference_answer"})  # type: ignore

    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
            parallel_workers: The number of processes to run in parallel

        Returns:
            Evaluation results
        """
        num_action: int = 0
        accuracies: list[float] = []
        run_results: list[AgentRunResult | None] = []

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

        assert self.dataset is not None
        dataset_to_use = self.dataset
        if evaluation_config.test_mode:
            dataset_to_use = [
                {
                    "question": "Open browser at url huggingface.co, then return 'opened browser'.",
                    "reference_answer": [
                        {"content": {"url": "https://huggingface.co/"}, "match_function_name": "url_exact_match"}
                    ],
                }
            ]

        date = datetime.date.today().isoformat()

        run_results = get_agent_responses(
            dataset_to_use,
            model,
            evaluation_config=evaluation_config,
            run_storage_folder=f"./output/mind2web/{evaluation_config.run_name}",
            reformulate_responses=False,
            return_browser_url=True,
        )

        # Evaluate all responses
        accuracies = []
        for result in run_results:  # type: ignore
            score = 0.0
            if result is None or not isinstance(result.answer, str):
                continue

            for element in result.reference_answer:
                # We'll average the scores across different evaluation criteria for this item
                if element["match_function_name"] == "url_included_match":
                    score += self.url_include_match(result.answer, element["content"]["url"])
                elif element["match_function_name"] == "url_exact_match":
                    score += self.url_exact_match(result.answer, element["content"]["url"])
                else:
                    continue
            score = float(score) / max(
                len(
                    [
                        item
                        for item in result.reference_answer
                        if item["match_function_name"] in ["url_included_match", "url_exact_match"]
                    ]
                ),
                1,
            )
            accuracies.append(score)
            num_action += 1

        if evaluation_config.test_mode:
            accuracies = [run_results[0].answer == "https://huggingface.co/"]  # type: ignore

        avg_accuracy = float(np.mean(accuracies)) if accuracies else 0.0

        evaluation_results = {
            "avg_accuracy": avg_accuracy,
            "proportion_missing": self._calculate_proportion_missing(run_results),
            "count_samples": len(run_results),
        }

        return BenchmarkResult(evaluation_results, "avg_accuracy")


mind2web_live = Mind2WebBenchmark(
    name="mind2web_live",
    config=Mind2WebConfig.create(),
    tags=["mind2web", "multistep", "hf_dataset", "online", "to_evaluate"],
)

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

    # model = OpenAIServerModel(
    #     model_id="gpt-4o",
    # )

    from smolagents import InferenceClientModel

    model = InferenceClientModel(
        model_id="https://a1ktmpegjfwv9n2y.us-east-1.aws.endpoints.huggingface.cloud",
        max_tokens=4096,
    )
    score = mind2web_live.evaluate(
        model=model, evaluation_config=EvaluationConfig(parallel_workers=1, run_name="test_run", test_mode=True)
    )
    logger.info(f"Score: {score}")
else:
    registry.register(mind2web_live)
