"""
xbench-DeepSearch dataset implementation.

This module provides a class for the xbench-DeepSearch benchmark dataset,
which evaluates deep research and search capabilities.
"""

import base64
from typing import Any, Dict, List
from loguru import logger
from .base import BenchmarkDataset


class XBenchDeepSearchDataset(BenchmarkDataset):
    """xbench-DeepSearch benchmark dataset for deep research evaluation."""

    @staticmethod
    def xor_decrypt(data: bytes, key: str) -> bytes:
        """XOR decrypt data with a key."""
        key_bytes = key.encode("utf-8")
        key_length = len(key_bytes)
        return bytes(
            [data[i] ^ key_bytes[i % key_length] for i in range(len(data))]
        )

    @classmethod
    def get_dataset_info(cls) -> Dict[str, str]:
        """Get basic information about the dataset."""
        return {
            "id": "xbench_deepsearch",
            "name": "xbench-DeepSearch",
            "description": "Deep research and search capability evaluation (100 questions)",
            "url": "https://huggingface.co/datasets/xbench/DeepSearch",
        }

    @classmethod
    def get_default_dataset_path(cls) -> str:
        """Get the default path for the dataset."""
        return "xbench/DeepSearch"  # Hugging Face dataset identifier

    def load(self) -> List[Dict[str, Any]]:
        """Override load to handle HuggingFace datasets directly.

        Keeps the base-class contract: callers configure sampling via the
        constructor (``num_examples``/``seed``) and call ``load()`` with no
        arguments, as ``DatasetRegistry.load_dataset()`` does.

        Returns:
            List of processed dataset examples
        """
        import random

        if self._is_loaded:
            return self.examples

        # Load the data
        data = self.load_data(self.dataset_path)

        # Sample if requested
        if self.num_examples and len(data) > self.num_examples:
            # Security: seeded random for reproducible benchmark sampling, not security-sensitive
            random.seed(self.seed)
            data = random.sample(data, self.num_examples)

        # Process each example
        self.examples = [self.process_example(item) for item in data]
        self._is_loaded = True
        return self.examples

    def load_data(
        self,
        dataset_path: str = None,
    ) -> List[Dict[str, Any]]:
        """Load the full xbench-DeepSearch dataset from Hugging Face.

        Sampling happens in load(), driven by the constructor's
        num_examples/seed.

        Args:
            dataset_path: Path to dataset (defaults to Hugging Face)

        Returns:
            List of questions from xbench-DeepSearch
        """
        try:
            from datasets import load_dataset
        except ImportError:
            logger.exception(
                "datasets library not installed. Run: pip install datasets"
            )
            # Fallback to direct download
            return self._load_from_url()

        dataset_path = dataset_path or self.get_default_dataset_path()

        try:
            logger.info(
                f"Loading xbench-DeepSearch dataset from {dataset_path}"
            )

            # Load the dataset from Hugging Face (no authentication needed)
            dataset = load_dataset(dataset_path, split="train")

            # Format for our benchmark system and decrypt
            formatted_questions = []
            for item in dataset:
                # Get the canary key for decryption
                canary = item.get("canary", "")

                # Decrypt prompt and answer if they're encrypted
                prompt = item.get("prompt", "")
                answer = item.get("answer", "")

                try:
                    # Try to decrypt if it looks like base64
                    if prompt and all(
                        c
                        in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
                        for c in prompt[:100]
                    ):
                        decrypted_prompt = self.xor_decrypt(
                            base64.b64decode(prompt), canary
                        ).decode("utf-8")
                        prompt = decrypted_prompt

                    if answer and all(
                        c
                        in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
                        for c in answer[:100]
                    ):
                        decrypted_answer = self.xor_decrypt(
                            base64.b64decode(answer), canary
                        ).decode("utf-8")
                        answer = decrypted_answer
                except Exception:
                    logger.warning("Failed to decrypt item")

                formatted_item = {
                    "id": item.get("id", f"xbench_{len(formatted_questions)}"),
                    "problem": prompt,
                    "answer": answer,
                    "reference_steps": item.get("reference_steps", ""),
                    "canary": canary,
                }
                formatted_questions.append(formatted_item)

            logger.info(
                f"Loaded {len(formatted_questions)} questions from xbench-DeepSearch"
            )
            return formatted_questions

        except Exception:
            logger.warning("Failed to load via datasets library")
            logger.info("Falling back to direct download")
            return self._load_from_url()

    def _load_from_url(self) -> List[Dict[str, Any]]:
        """Load the full dataset directly from URL without datasets library.

        Returns:
            List of questions from xbench-DeepSearch
        """
        import pandas as pd

        try:
            # Direct URL to the CSV file on Hugging Face
            url = "https://huggingface.co/datasets/xbench/DeepSearch/resolve/main/data/train-00000-of-00001.parquet"

            logger.info(f"Downloading xbench-DeepSearch from {url}")
            df = pd.read_parquet(url)

            # Convert to list of dicts and decrypt
            questions = []
            for _, row in df.iterrows():
                # Get the canary key for decryption
                canary = row.get("canary", "")

                # Decrypt prompt and answer
                prompt = row.get("prompt", "")
                answer = row.get("answer", "")

                try:
                    if prompt and all(
                        c
                        in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
                        for c in prompt[:100]
                    ):
                        decrypted_prompt = self.xor_decrypt(
                            base64.b64decode(prompt), canary
                        ).decode("utf-8")
                        prompt = decrypted_prompt

                    if answer and all(
                        c
                        in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
                        for c in answer[:100]
                    ):
                        decrypted_answer = self.xor_decrypt(
                            base64.b64decode(answer), canary
                        ).decode("utf-8")
                        answer = decrypted_answer
                except Exception:
                    logger.warning("Failed to decrypt item")

                questions.append(
                    {
                        "id": row.get("id", f"xbench_{len(questions)}"),
                        "problem": prompt,
                        "answer": answer,
                        "reference_steps": row.get("reference_steps", ""),
                        "canary": canary,
                    }
                )

            logger.info(
                f"Loaded {len(questions)} questions via direct download"
            )
            return questions

        except Exception:
            logger.exception("Failed to load dataset")
            return []

    def process_example(self, example: Dict[str, Any]) -> Dict[str, Any]:
        """Process a single example from the dataset.

        xbench-DeepSearch questions are designed for deep research evaluation.
        """
        processed = dict(example)

        # Add evaluation metadata
        processed["requires_deep_search"] = True
        processed["expected_iterations"] = (
            4  # Deep search questions need multiple iterations
        )

        # Evaluation criteria for research questions
        processed["evaluation_criteria"] = {
            "accuracy": 0.4,
            "completeness": 0.3,
            "reasoning": 0.2,
            "sources": 0.1,  # Credit for citing sources
        }

        return processed
