"""
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, Optional
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,
        dataset_path: str = None,
        num_examples: Optional[int] = None,
        seed: int = 42,
    ) -> List[Dict[str, Any]]:
        """Override load to handle HuggingFace datasets directly.

        Args:
            dataset_path: Path to dataset (defaults to HuggingFace)
            num_examples: Optional number of examples to limit
            seed: Random seed for sampling

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

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

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

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

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

        Args:
            dataset_path: Path to dataset (defaults to Hugging Face)
            num_examples: Optional number of examples to limit
            seed: Random seed for sampling

        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(num_examples, seed)

        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 as e:
                    logger.warning(f"Failed to decrypt item: {e}")

                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)

            # Apply sampling if num_examples is specified
            if num_examples and num_examples < len(formatted_questions):
                import random

                # Use provided seed or None for random sampling
                # Security: seeded random for reproducible benchmark sampling, not security-sensitive
                if seed is not None:
                    random.seed(seed)
                formatted_questions = random.sample(
                    formatted_questions, num_examples
                )
                logger.info(
                    f"Sampled {num_examples} questions from xbench-DeepSearch (total available: {len(dataset)})"
                )
            else:
                logger.info(
                    f"Loaded {len(formatted_questions)} questions from xbench-DeepSearch"
                )
            return formatted_questions

        except Exception as e:
            logger.warning(f"Failed to load via datasets library: {e}")
            logger.info("Falling back to direct download")
            return self._load_from_url(num_examples, seed)

    def _load_from_url(
        self, num_examples: int = None, seed: int = None
    ) -> List[Dict[str, Any]]:
        """Load dataset directly from URL without datasets library.

        Args:
            num_examples: Optional number of examples to limit
            seed: Random seed for sampling

        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 as e:
                    logger.warning(f"Failed to decrypt item: {e}")

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

            # Apply sampling if num_examples is specified
            if num_examples and num_examples < len(questions):
                import random

                # Use provided seed or None for random sampling
                # Security: seeded random for reproducible benchmark sampling, not security-sensitive
                if seed is not None:
                    random.seed(seed)
                questions = random.sample(questions, num_examples)
                logger.info(
                    f"Sampled {num_examples} questions via direct download (total available: {len(questions)})"
                )
            else:
                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
