"""
Mixed Benchmark Sampler for tag-based random sampling across multiple benchmarks.

This module enables training and evaluation on random samples from multiple benchmarks
that share common tags (e.g., 'coding', 'reasoning', 'math').
"""

import random
import logging
from typing import List, Dict, Any, Optional, Set, Tuple
from dataclasses import dataclass
from collections import defaultdict

from wisent.core.utils.infra_tools.errors import UnknownTypeError, InsufficientDataError

# Suppress BigCode debug output
import builtins
_original_print = getattr(builtins, '_original_print', builtins.print)

def _quiet_print(*args, **kwargs):
    """Filter out BigCode debug messages."""
    message = ' '.join(str(arg) for arg in args)
    if any(x in message for x in ['DEBUG', 'Available tasks:', 'ERROR extracting', 'bigcode_eval']):
        return
    _original_print(*args, **kwargs)

# Store original print and patch
builtins._original_print = builtins.print
builtins.print = _quiet_print

from wisent.core.primitives.models.lm_harness_integration.only_benchmarks import CORE_BENCHMARKS
from wisent.core.primitives.contrastive_pairs import ContrastivePairSet
from .cache.managed_cached_benchmarks import ManagedCachedBenchmarks, get_managed_cache
from wisent.core.utils.config_tools.constants import DISPLAY_TOP_N_SMALL

logger = logging.getLogger(__name__)


@dataclass
class BenchmarkSample:
    """A single sample from a benchmark."""
    benchmark_name: str
    sample_data: Dict[str, Any]
    tags: List[str]


class MixedBenchmarkSampler:
    """
    Samples randomly from multiple benchmarks based on tags.
    
    This creates more robust classifiers by training on diverse data
    from multiple sources rather than a single benchmark.
    """
    
    def __init__(self, chunk_size: int, max_cache_age_days: int, cache_dir: str = "./benchmark_cache"):
        """
        Initialize the mixed benchmark sampler.

        Args:
            chunk_size: Number of samples per cache chunk
            max_cache_age_days: Maximum age of cached data in days
            cache_dir: Directory for cached benchmark data
        """
        self.cache_dir = cache_dir
        self.managed_cache = get_managed_cache(chunk_size=chunk_size, max_cache_age_days=max_cache_age_days, cache_dir=cache_dir)
        self._benchmark_registry = self._build_benchmark_registry()
    
    def _build_benchmark_registry(self) -> Dict[str, List[str]]:
        """Build a registry mapping tags to benchmark names."""
        tag_to_benchmarks = defaultdict(list)
        
        for benchmark_name, config in CORE_BENCHMARKS.items():
            tags = config.get("tags", [])
            for tag in tags:
                tag_to_benchmarks[tag].append(benchmark_name)
        
        return dict(tag_to_benchmarks)
    
    def get_benchmarks_by_tag(self, tag: str) -> List[str]:
        """Get all benchmarks that have a specific tag."""
        return self._benchmark_registry.get(tag, [])
    
    def get_benchmarks_by_tags(self, tags: List[str], mode: str) -> List[str]:
        """
        Get benchmarks that match the given tags.
        
        Args:
            tags: List of tags to match
            mode: "any" (benchmark has at least one tag) or "all" (benchmark has all tags)
            
        Returns:
            List of benchmark names matching the criteria
        """
        if mode == "any":
            # Get benchmarks that have ANY of the specified tags
            matching_benchmarks = set()
            for tag in tags:
                matching_benchmarks.update(self.get_benchmarks_by_tag(tag))
            return list(matching_benchmarks)
        
        elif mode == "all":
            # Get benchmarks that have ALL of the specified tags
            if not tags:
                return []
            
            # Start with benchmarks that have the first tag
            matching_benchmarks = set(self.get_benchmarks_by_tag(tags[0]))
            
            # Intersect with benchmarks for each additional tag
            for tag in tags[1:]:
                matching_benchmarks &= set(self.get_benchmarks_by_tag(tag))
            
            return list(matching_benchmarks)
        
        else:
            raise UnknownTypeError(entity_type="mode", value=mode, valid_values=["any", "all"])
    
    def sample_mixed_dataset(
        self,
        tags: List[str],
        total_samples: int,
        tag_mode: str,
        split_ratio: float = 0.8,
        random_seed: Optional[int] = None,
        benchmark_weights: Optional[Dict[str, float]] = None, *, train_ratio: float,
    ) -> Tuple[List[BenchmarkSample], List[BenchmarkSample]]:
        """
        Sample a mixed dataset from benchmarks matching the given tags.

        Args:
            tags: Tags to filter benchmarks (e.g., ["coding", "python"])
            total_samples: Total number of samples to collect
            split_ratio: Train/test split ratio
            random_seed: Random seed for reproducibility
            tag_mode: "any" or "all" for tag matching
            benchmark_weights: Optional weights for sampling probability per benchmark

        Returns:
            Tuple of (train_samples, test_samples)
        """
        if random_seed is not None:
            random.seed(random_seed)
        
        # Get matching benchmarks
        matching_benchmarks = self.get_benchmarks_by_tags(tags, mode=tag_mode)
        
        if not matching_benchmarks:
            raise InsufficientDataError(reason=f"No benchmarks found with tags {tags} (mode={tag_mode})")
        
        logger.info(f"Found {len(matching_benchmarks)} benchmarks matching tags {tags}")
        logger.info(f"Matching benchmarks: {matching_benchmarks[:DISPLAY_TOP_N_SMALL]}...")
        
        # Collect all available samples from matching benchmarks
        all_samples = []
        benchmark_sample_counts = {}
        
        # Skip benchmarks that require code execution permission
        code_execution_benchmarks = {"apps", "ds1000", "mercury"}
        
        for benchmark_name in matching_benchmarks:
            # Skip benchmarks that require code execution for safety
            if benchmark_name in code_execution_benchmarks:
                logger.info(f"Skipping {benchmark_name} (requires code execution permission)")
                continue
                
            try:
                # Get samples from this benchmark
                samples_per_benchmark = max(DISPLAY_TOP_N_SMALL, total_samples // len(matching_benchmarks))
                
                cached_samples = self.managed_cache.get_task_samples(
                    task_name=benchmark_name,
                    limit=samples_per_benchmark,
                    force_fresh=False,
                    train_ratio=train_ratio,
                )
                
                # Convert to BenchmarkSample objects
                for sample in cached_samples:
                    benchmark_sample = BenchmarkSample(
                        benchmark_name=benchmark_name,
                        sample_data=sample,
                        tags=CORE_BENCHMARKS[benchmark_name].get("tags", [])
                    )
                    all_samples.append(benchmark_sample)
                
                benchmark_sample_counts[benchmark_name] = len(cached_samples)
                
            except Exception as e:
                logger.warning(f"Failed to load samples from {benchmark_name}: {e}")
                continue
        
        if not all_samples:
            raise InsufficientDataError(reason=f"No samples could be loaded from any benchmark with tags {tags}")
        
        logger.info(f"Collected {len(all_samples)} total samples from {len(benchmark_sample_counts)} benchmarks")
        for benchmark, count in benchmark_sample_counts.items():
            logger.debug(f"  {benchmark}: {count} samples")
        
        # Apply benchmark weights if provided
        if benchmark_weights:
            weighted_samples = []
            for sample in all_samples:
                weight = benchmark_weights[sample.benchmark_name]
                # Duplicate samples based on weight (simple approach)
                weighted_samples.extend([sample] * int(weight))
            all_samples = weighted_samples
        
        # Randomly sample and shuffle
        if len(all_samples) > total_samples:
            all_samples = random.sample(all_samples, total_samples)
        else:
            # If we have fewer samples than requested, use all and log warning
            logger.warning(f"Only {len(all_samples)} samples available, requested {total_samples}")
        
        random.shuffle(all_samples)
        
        # Split into train/test
        split_point = int(len(all_samples) * split_ratio)
        train_samples = all_samples[:split_point]
        test_samples = all_samples[split_point:]
        
        # Log distribution
        train_dist = defaultdict(int)
        test_dist = defaultdict(int)
        
        for sample in train_samples:
            train_dist[sample.benchmark_name] += 1
        
        for sample in test_samples:
            test_dist[sample.benchmark_name] += 1
        
        logger.info(f"Train set: {len(train_samples)} samples from {len(train_dist)} benchmarks")
        logger.info(f"Test set: {len(test_samples)} samples from {len(test_dist)} benchmarks")
        
        return train_samples, test_samples
    
    def extract_contrastive_pairs_from_mixed_samples(
        self,
        samples: List[BenchmarkSample]
    ) -> List[Dict[str, Any]]:
        """
        Extract contrastive pairs from mixed benchmark samples.
        
        Args:
            samples: List of BenchmarkSample objects
            
        Returns:
            List of contrastive pairs with question, correct_answer, incorrect_answer
        """
        contrastive_pairs = []
        
        for sample in samples:
            try:
                # Each sample already has normalized QA pair from managed cache
                qa_pair = sample.sample_data.get("normalized", {})
                
                if qa_pair and all(k in qa_pair for k in ["question", "correct_answer", "incorrect_answer"]):
                    # Add benchmark source info
                    qa_pair["source_benchmark"] = sample.benchmark_name
                    qa_pair["tags"] = sample.tags
                    contrastive_pairs.append(qa_pair)
                else:
                    logger.warning(f"Invalid QA pair from {sample.benchmark_name}")
                    
            except Exception as e:
                logger.warning(f"Failed to extract pair from {sample.benchmark_name}: {e}")
                continue
        
        logger.info(f"Extracted {len(contrastive_pairs)} contrastive pairs from mixed samples")
        
        return contrastive_pairs

    def create_mixed_contrastive_pair_set(self, tags, total_samples, name=None, **kwargs):
        """Create a ContrastivePairSet from mixed benchmark samples."""
        from wisent.core.utils.services.benchmarks._mixed_benchmark_helpers import create_mixed_contrastive_pair_set_impl
        return create_mixed_contrastive_pair_set_impl(self, tags, total_samples, name, **kwargs)


# Re-export for backward compatibility
from wisent.core.utils.services.benchmarks._mixed_benchmark_helpers import sample_benchmarks_by_tag  # noqa: E402
