import bisect

# import random
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

from .checkpointable_prng import CheckpointablePRNG


@dataclass
class HistogramBucket:
    """Pure mathematical histogram bucket for InverseTransformSampler"""

    lower_bound: float
    upper_bound: float
    count: int  # Number of values in THIS bucket, not cumulative now. ATTENTION, in TiDB's doc, count is cumulative but in real db it's not.
    repeats: int = 0  # Number of values equal to upper_bound
    ndv: int = 0  # Specially in TiDB, it is always 0. And ndv is not used in sampling.
    bucket_id: Optional[int] = None  # Original bucket ID from stats_buckets table
    column_name: Optional[str] = None  # Column name for identification
    data_type: Optional[str] = None  # Underlying column data type (e.g., INT, DATETIME)
    raw_lower_bound: Optional[str] = (
        None  # Original lower bound text for precision-sensitive domains
    )
    raw_upper_bound: Optional[str] = (
        None  # Original upper bound text for precision-sensitive domains
    )


@dataclass
class SampleResult:
    """Result of sampling with source type information"""

    value: Optional[float]
    source_type: str  # "null", "topn", "bucket"
    bucket_id: Optional[int] = None  # Original bucket ID if from bucket
    column_name: Optional[str] = None  # Column name if from bucket


@dataclass
class TopNItem:
    """Top N item with value and count"""

    value: float
    count: int


class InverseTransformSampler:
    """
    Optimized mathematical inverse transform sampler for efficient row-by-row sampling
    Supports multiple input formats for maximum flexibility:
    - TiDB format: buckets + top_n with cumulative counts (backward compatible)
    - PMF format: values + probabilities (most general)
    - Histogram format: bins + counts (common statistical format)
    - Legacy tuple format: for backward compatibility with prepare.py
    - Absolute count format: standard histogram buckets with per-bucket counts

    Handles both cumulative counts (TiDB style) and absolute counts (standard histogram style)
    via the is_cumulative_count parameter.
    """

    def __init__(
        self,
        distribution_data: Optional[Dict] = None,
        buckets: Optional[List] = None,  # Can be HistogramBucket objects or tuples
        top_n: Optional[List] = None,  # Can be TopNItem objects or tuples
        null_count: int = 0,
        is_cumulative_count: bool = True,  # True for TiDB format, False for absolute counts
        seed: Optional[int] = 42,
    ):
        # Create independent random state for this sampler instance
        self._random = CheckpointablePRNG(seed)

        self.null_count = null_count

        # Parse input data - priority: distribution_data > buckets/top_n
        if distribution_data is not None:
            self.buckets, self.top_n = self._parse_distribution_data(distribution_data)
        else:
            assert (
                buckets is not None or top_n is not None
            ), "Must provide either distribution_data or buckets/top_n"
            normalized_buckets = self._normalize_buckets(buckets or [])
            # Convert cumulative counts to absolute counts if needed
            if is_cumulative_count and normalized_buckets:
                self.buckets = self._convert_cumulative_buckets(normalized_buckets)
            else:
                self.buckets = normalized_buckets
            self.top_n = self._normalize_top_n(top_n or [])

        # Build optimized CDF for efficient sampling
        self._build_cdf()

    def _normalize_buckets(self, buckets_input):
        """Convert various bucket formats to HistogramBucket objects"""
        normalized_buckets = []

        for bucket in buckets_input:
            if isinstance(bucket, HistogramBucket):
                # Already correct format
                normalized_buckets.append(bucket)
            else:
                raise ValueError(f"Unsupported bucket format: {bucket}")

        return normalized_buckets

    def _normalize_top_n(self, top_n_input):
        """Convert various TopN formats to TopNItem objects"""
        normalized_top_n = []

        for item in top_n_input:
            if isinstance(item, TopNItem):
                # Already correct format
                normalized_top_n.append(item)
            else:
                raise ValueError(f"Unsupported TopN format: {item}")

        return normalized_top_n

    def _convert_cumulative_buckets(self, buckets):
        """
        Convert TiDB cumulative bucket counts to absolute counts

        TiDB stores count as cumulative: each bucket's count includes all previous buckets.
        This method converts them to absolute counts per bucket.

        Args:
            buckets: List of HistogramBucket objects with cumulative counts

        Returns:
            List of HistogramBucket objects with absolute counts
        """
        if not buckets:
            return []

        # Sort buckets to ensure correct order for cumulative processing
        # Assuming buckets are ordered by lower_bound, but let's be safe
        sorted_buckets = sorted(buckets, key=lambda b: b.lower_bound)

        converted_buckets = []
        prev_cumulative = 0

        for bucket in sorted_buckets:
            current_cumulative = bucket.count
            actual_count = current_cumulative - prev_cumulative

            # Only keep buckets with positive count
            if actual_count > 0:
                converted_buckets.append(
                    HistogramBucket(
                        lower_bound=bucket.lower_bound,
                        upper_bound=bucket.upper_bound,
                        count=actual_count,
                        repeats=bucket.repeats,
                        bucket_id=bucket.bucket_id,
                        column_name=bucket.column_name,
                        data_type=bucket.data_type,
                        raw_lower_bound=bucket.raw_lower_bound,
                        raw_upper_bound=bucket.raw_upper_bound,
                    )
                )

            prev_cumulative = current_cumulative

        return converted_buckets

    def _parse_distribution_data(
        self, data: Dict
    ) -> Tuple[List[HistogramBucket], List[TopNItem]]:
        """
        Parse various distribution data formats into buckets and top_n lists

        Supported formats:
        1. PMF (Probability Mass Function) - most general:
           {
               "type": "pmf",
               "values": [1.0, 2.0, 3.0],
               "probabilities": [0.5, 0.3, 0.2]
           }

        2. Histogram format:
           {
               "type": "histogram",
               "bins": [0, 10, 20, 30],  # bin edges
               "counts": [100, 200, 150]  # counts per bin
           }

        Args:
            data: Dictionary containing distribution specification

        Returns:
            tuple: (buckets, top_n) lists compatible with existing logic
        """
        if not isinstance(data, dict) or "type" not in data:
            raise ValueError("distribution_data must be a dict with 'type' field")

        dist_type = data["type"].lower()

        if dist_type == "pmf":
            return self._parse_pmf_format(data)
        elif dist_type == "histogram":
            return self._parse_histogram_format(data)
        else:
            raise ValueError(
                f"Unsupported distribution type: {dist_type}. "
                f"Supported types: 'pmf', 'histogram'"
            )

    def _parse_pmf_format(
        self, data: Dict
    ) -> Tuple[List[HistogramBucket], List[TopNItem]]:
        """
        Parse PMF (Probability Mass Function) format

        PMF is the most general discrete distribution representation:
        Each value has an exact probability of occurrence
        """
        required_fields = ["values", "probabilities"]
        for field in required_fields:
            if field not in data:
                raise ValueError(f"PMF format requires '{field}' field")

        values = data["values"]
        probabilities = data["probabilities"]

        if len(values) != len(probabilities):
            raise ValueError(
                "PMF format: values and probabilities must have same length"
            )

        if not values:
            return [], []

        # Convert probabilities to counts (scale by a reasonable factor)
        # This maintains relative probabilities while giving integer counts
        prob_sum = sum(probabilities)
        if abs(prob_sum - 1.0) > 1e-6:
            raise ValueError(f"PMF probabilities must sum to 1.0, got {prob_sum}")

        # Scale to reasonable integer counts (using 10000 as base)
        scale_factor = 10000
        counts = [max(1, int(prob * scale_factor)) for prob in probabilities]

        # Create TopNItem entries for all PMF values (they're all "exact" values)
        top_n = []
        for value, count in zip(values, counts):
            top_n.append(TopNItem(value=float(value), count=count))

        # No buckets needed for PMF (all values are exact)
        buckets = []

        return buckets, top_n

    def _parse_histogram_format(
        self, data: Dict
    ) -> Tuple[List[HistogramBucket], List[TopNItem]]:
        """
        Parse histogram format with bins and counts

        This represents continuous distributions with binned data
        """
        required_fields = ["bins", "counts"]
        for field in required_fields:
            if field not in data:
                raise ValueError(f"Histogram format requires '{field}' field")

        bins = data["bins"]
        counts = data["counts"]

        if not counts:
            return [], []

        if len(bins) != len(counts) + 1:
            raise ValueError(
                "Histogram format: bins must have one more element than counts"
            )

        # Create HistogramBucket entries
        buckets = []
        for i, count in enumerate(counts):
            if count > 0:  # Skip empty bins
                lower_bound = float(bins[i])
                upper_bound = float(bins[i + 1])

                # For histogram bins, assume no specific "repeats" at upper bound
                # This could be enhanced if repeat information is provided
                repeats = 0

                buckets.append(
                    HistogramBucket(
                        lower_bound=lower_bound,
                        upper_bound=upper_bound,
                        count=int(count),
                        repeats=repeats,
                    )
                )

        # No TopN items for pure histogram (could be enhanced to detect peaks)
        top_n = []

        return buckets, top_n

    def _build_cdf(self):
        """
        Build optimized cumulative distribution function for fast sampling
        It is not a real CDF or inverse CDF, but a pre-computed structure
        """
        # Calculate total count
        topn_count = sum(item.count for item in self.top_n)
        bucket_count = sum(bucket.count for bucket in self.buckets)
        self.total_count = topn_count + bucket_count + self.null_count

        if self.total_count == 0:
            self.cdf_probs = []
            self.cdf_values = []
            self.cdf_types = []
            self.cdf_buckets = []
            self.null_probability = 1.0
            return

        # Calculate null probability
        self.null_probability = self.null_count / self.total_count

        # Build optimized parallel arrays for fastest lookup
        self.cdf_probs = []  # Cumulative probabilities for binary search
        self.cdf_values = []  # Pre-computed values for TopN items
        self.cdf_types = []  # Entry type: 0=topn, 1=bucket
        self.cdf_buckets = []  # Bucket objects for bucket entries

        # Pre-compute repeat probabilities for buckets to avoid division in sample()
        self.bucket_repeat_probs = {}
        for bucket in self.buckets:
            if bucket.repeats > 0:
                self.bucket_repeat_probs[id(bucket)] = bucket.repeats / bucket.count

        cumulative = self.null_probability

        # Add TopN entries (sorted by value for consistency)
        sorted_topn = sorted(self.top_n, key=lambda x: x.value)

        for item in sorted_topn:
            prob = item.count / self.total_count
            cumulative += prob
            self.cdf_probs.append(cumulative)
            self.cdf_values.append(item.value)
            self.cdf_types.append(0)  # TopN type
            self.cdf_buckets.append(None)

        # Add Bucket entries
        for bucket in self.buckets:
            prob = bucket.count / self.total_count
            cumulative += prob
            self.cdf_probs.append(cumulative)
            self.cdf_values.append(None)  # No pre-computed value for buckets
            self.cdf_types.append(1)  # Bucket type
            self.cdf_buckets.append(bucket)

    def sample(self) -> SampleResult:
        """
        Sample a numeric value using optimized inverse transform method
        Returns detailed result with source information

        Returns:
            SampleResult: Sampled result with value and source information
        """
        if self.total_count == 0:
            return SampleResult(value=None, source_type="null")

        u = self._random.random()

        # Fast null check
        if u < self.null_probability:
            return SampleResult(value=None, source_type="null")

        # Fast binary search in pre-computed probabilities
        idx = bisect.bisect_left(self.cdf_probs, u)
        if idx >= len(self.cdf_probs):
            idx = len(self.cdf_probs) - 1

        # Fast dispatch based on pre-computed type
        if self.cdf_types[idx] == 0:  # TopN - direct value lookup
            return SampleResult(value=self.cdf_values[idx], source_type="topn")
        else:  # Bucket - minimal computation sampling
            bucket = self.cdf_buckets[idx]

            # Fast repeat handling using pre-computed probability
            bucket_id = id(bucket)
            if bucket_id in self.bucket_repeat_probs:
                if self._random.random() < self.bucket_repeat_probs[bucket_id]:
                    # Return upper_bound but still mark as bucket (business layer knows to use upper_bound)
                    return SampleResult(
                        value=bucket.upper_bound,
                        source_type="bucket",
                        bucket_id=bucket.bucket_id,
                        column_name=bucket.column_name,
                    )

            # Direct uniform sampling
            sampled_value = self._random.uniform(bucket.lower_bound, bucket.upper_bound)
            return SampleResult(
                value=sampled_value,
                source_type="bucket",
                bucket_id=bucket.bucket_id,
                column_name=bucket.column_name,
            )

    def sample_batch(self, n: int) -> List[SampleResult]:
        """
        Sample multiple values

        Args:
            n: Number of samples to generate

        Returns:
            List of sampled results
        """
        return [self.sample() for _ in range(n)]

    def get_statistics(self) -> Dict:
        """
        Get statistics about the sampler

        Returns:
            dict: Statistics including bucket info, topn info, etc.
        """
        return {
            "num_buckets": len(self.buckets),
            "num_topn": len(self.top_n),
            "total_count": self.total_count,
            "null_count": self.null_count,
            "null_probability": self.null_probability,
            "bucket_ranges": [
                (b.lower_bound, b.upper_bound, b.count) for b in self.buckets
            ],
            "topn_values": [(t.value, t.count) for t in self.top_n],
            "optimization_info": {
                "cdf_entries": len(self.cdf_probs) if hasattr(self, "cdf_probs") else 0,
                "pre_computed_repeat_probs": (
                    len(self.bucket_repeat_probs)
                    if hasattr(self, "bucket_repeat_probs")
                    else 0
                ),
            },
        }