"""
Adapted from BrowseComp: A Simple Yet Challenging Benchmark for Browsing Agents
Authors: Jason Wei, Zhiqing Sun, Spencer Papay, Scott McKinney, Jeffrey Han, Isa Fulford, Hyung Won Chung, Alex Tachard Passos, William Fedus, Mia Glaese
https://openai.com/index/browsecomp/
"""

import base64
import hashlib
import re

import numpy as np
from smolagents import InferenceClientModel, Model

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmark
from screensuite.benchmarks.multistep.browse_comp.config import BrowseCompConfig
from screensuite.benchmarks.multistep.browse_comp.prompts import (
    GRADER_TEMPLATE,
    QUERY_TEMPLATE,
)
from screensuite.benchmarks.multistep.config import AgentRunResult
from screensuite.benchmarks.multistep.generation import get_agent_responses
from screensuite.registry import registry


def derive_key(password: str, length: int) -> bytes:
    """Derive a fixed-length key from the password using SHA256."""
    hasher = hashlib.sha256()
    hasher.update(password.encode())
    key = hasher.digest()
    return key * (length // len(key)) + key[: length % len(key)]


def decrypt(ciphertext_b64: str, password: str) -> str:
    """Decrypt base64-encoded ciphertext with XOR."""
    encrypted = base64.b64decode(ciphertext_b64)
    key = derive_key(password, len(encrypted))
    decrypted = bytes(a ^ b for a, b in zip(encrypted, key))
    return decrypted.decode()


class BrowseCompBenchmark(HubBaseBenchmark[BrowseCompConfig, None]):
    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({"problem": "question", "answer": "reference_answer"})  # type: ignore

        def process_row(row):
            processed_row = row.copy()
            processed_row["question"] = QUERY_TEMPLATE.format(Question=decrypt(row["question"], row["canary"]))
            processed_row["reference_answer"] = decrypt(row["reference_answer"], row["canary"])
            return processed_row

        self.dataset = self.dataset.map(process_row)

    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
        """
        self.grader_model = InferenceClientModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", provider="together")
        accuracies: list[float] = []
        run_results: list[AgentRunResult] = []

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

        run_results = get_agent_responses(
            self.dataset,  # type: ignore
            model,
            evaluation_config=evaluation_config,
            run_storage_folder=f"./output/browse_comp/{evaluation_config.run_name}",
            reformulate_responses=False,
            return_browser_url=False,
        )

        # Evaluate all responses
        for result in run_results:
            if result is None:
                continue

            grader_prompt = GRADER_TEMPLATE.format(
                question=result.question,
                correct_answer=result.reference_answer,
                response=result.answer,
            )

            prompt_messages = [dict(content=grader_prompt, role="user")]
            # Handler grader model errors
            try:
                grading_response = self.grader_model.generate(prompt_messages).content
            except Exception as e:
                print(f"Error grading response: {e}")
                grading_response = "no"
            finally:
                if grading_response is None:
                    grading_response = "no"

            match = re.search(r"correct: (yes|no)", grading_response)
            grade_result = match.group(0).split(": ")[-1] if match else "no"  # Default to "no" if no

            is_correct = grade_result == "yes"
            score = is_correct

            accuracies.append(score)

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


browse_comp = BrowseCompBenchmark(
    name="browse_comp",
    config=BrowseCompConfig(),
    tags=["browse_comp", "multistep", "hf_dataset", "online", "to_evaluate"],
)

registry.register(browse_comp)
