import math
from decimal import ROUND_CEILING, ROUND_FLOOR, Decimal, InvalidOperation
from typing import List, Optional

from .block_shuffler import BlockShuffler
from .checkpointable_prng import CheckpointablePRNG
from .inverse_transform_sampler import HistogramBucket, InverseTransformSampler
from .sampling_phase import SamplingPhase
from .tidb_types import TiDBNumericBucket, TiDBNumericTopNItem


class IntSampler:
    """
    Integer sampler using InverseTransformSampler
    Handles TiDB integer column statistics
    """

    def __init__(
        self,
        buckets: List[
            TiDBNumericBucket
        ],  # TiDB numeric bucket objects from schema_manager
        top_n: Optional[
            List[TiDBNumericTopNItem]
        ] = None,  # TiDB numeric TopN objects from schema_manager
        null_count: int = 0,
        is_cumulative_count: bool = False,  # False for absolute counts
        constraint_type: str = "normal",  # "normal" or "unique"
        seed: Optional[int] = 42,
        column_name: str = None,  # Column name for logging
    ):
        """
        Initialize IntSampler from standardized TiDB numeric objects

        Args:
            buckets: List of TiDB numeric bucket objects (standardized by schema_manager)
            top_n: List of TiDB numeric TopN objects (standardized by schema_manager)
            null_count: Number of null values
            is_cumulative_count: False for absolute counts
            constraint_type: "normal" for InverseTransformSampler, "unique" for BlockShuffler
            seed: Random seed
        """
        self.constraint_type = constraint_type
        self.null_count = null_count
        self.column_name = column_name or "unknown_column"
        self._seed = seed

        # Track IntSampler instances for i_id
        self._null_rng = None
        if null_count > 0:
            self._null_rng = CheckpointablePRNG(seed=seed)

        # Convert TiDB numeric objects to mathematical objects
        math_buckets = [bucket.to_histogram_bucket() for bucket in (buckets or [])]
        math_topn = [item.to_topn_item() for item in (top_n or [])]

        # Choose underlying sampler based on constraint type
        if constraint_type == "unique":
            if math_buckets:
                self._clamp_unique_bucket_counts(math_buckets)
            # Primary key/unique key uses BlockShuffler (no duplicate sampling)
            # Handle case where column has no statistical data
            if not math_buckets and not math_topn:
                # print(f"INFO: Unique integer column '{self.column_name}' has no statistical data, using pure new value generation")
                self.sampler = None  # No BlockShuffler needed
            else:
                self.sampler = BlockShuffler(
                    buckets=math_buckets,
                    top_n=math_topn,
                    seed=seed,
                )
            # Initialize new value generation strategy for unique columns
            self._init_new_value_generation(math_buckets, math_topn)
        else:
            # Normal column uses InverseTransformSampler (probability sampling)
            self.sampler = InverseTransformSampler(
                buckets=math_buckets,
                top_n=math_topn,
                null_count=null_count,
                is_cumulative_count=is_cumulative_count,
                seed=seed,
            )

    @staticmethod
    def _clamp_unique_bucket_counts(math_buckets: List[HistogramBucket]) -> None:
        """Clamp unique integer bucket counts to available integer span."""

        for bucket in math_buckets:
            if bucket.repeats != 1:
                continue

            lower_int = None
            upper_int = None

            lower_raw = getattr(bucket, "raw_lower_bound", None)
            upper_raw = getattr(bucket, "raw_upper_bound", None)

            if lower_raw is not None and upper_raw is not None:
                try:
                    lower_dec = Decimal(str(lower_raw))
                    upper_dec = Decimal(str(upper_raw))
                    lower_int = int(lower_dec.to_integral_value(rounding=ROUND_CEILING))
                    upper_int = int(upper_dec.to_integral_value(rounding=ROUND_FLOOR))
                except (InvalidOperation, ValueError):
                    lower_int = None
                    upper_int = None

            if lower_int is None or upper_int is None:
                lower_int = math.ceil(bucket.lower_bound)
                upper_int = math.floor(bucket.upper_bound)

            if upper_int < lower_int:
                capacity = 0
            else:
                capacity = upper_int - lower_int + 1

            if capacity < bucket.count:
                bucket.count = capacity

    @property
    def its(self):
        """Backward compatibility: Access underlying InverseTransformSampler"""
        if self.constraint_type == "unique":
            # BlockShuffler mode, return None or raise error for backward compatibility
            raise AttributeError(
                "IntSampler in 'unique' constraint mode uses BlockShuffler, "
                "not InverseTransformSampler. Use sampler.sampler instead."
            )
        return self.sampler

    def _init_new_value_generation(self, math_buckets, math_topn):
        """Initialize new value generation strategy for unique columns when BlockShuffler is exhausted"""
        # Analyze value range from statistical information
        all_values = []

        # Collect values from TopN
        for topn in math_topn:
            all_values.append(topn.value)

        # Collect values from buckets (use bounds as reference)
        for bucket in math_buckets:
            all_values.extend([bucket.lower_bound, bucket.upper_bound])

        if all_values:
            self._min_value = int(min(all_values))
            self._max_value = int(max(all_values))
            self._next_new_value = self._max_value + 1
        else:
            # No statistical data, start from 1
            self._min_value = 1
            self._max_value = 1
            self._next_new_value = 1

        # Flag to log exhaustion message only once
        self._block_exhausted_logged = False

    def _generate_new_unique_value(self) -> int:
        """Generate new unique value using incremental strategy"""
        value = self._next_new_value
        self._next_new_value += 1
        return value

    def has_statistical_capacity(self) -> bool:
        """Check if statistical information still has capacity"""
        if self.constraint_type == "unique":
            if self.sampler is None:
                return False  # No statistical data, no capacity
            return hasattr(self.sampler, "has_capacity") and self.sampler.has_capacity()
        return True  # Normal columns always have capacity

    def sample_with_phase(self, phase: SamplingPhase) -> Optional[int]:
        """Sample with phase awareness"""
        if self.constraint_type == "normal":
            # Normal columns use InverseTransformSampler in both phases
            sample_result = self.sampler.sample()
            return (
                int(round(sample_result.value))
                if sample_result.value is not None
                else None
            )

        elif self.constraint_type == "unique":
            if phase == SamplingPhase.PREPARE:
                # PREPARE phase: BlockShuffler + new value generation
                if self.sampler is None:
                    # No statistical data, use pure new value generation
                    return self._generate_new_unique_value()

                sample_result = self.sampler.sample()
                if sample_result is None:
                    # BlockShuffler exhausted, log once and switch to new value generation
                    if not self._block_exhausted_logged:
                        # remaining_capacity = getattr(
                        #     self.sampler, "remaining_capacity", "unknown"
                        # )
                        # print(
                        #     f"INFO: BlockShuffler exhausted for unique integer column '{self.column_name}', switching to new value generation. Remaining capacity: {remaining_capacity}"
                        # )
                        self._block_exhausted_logged = True
                    return self._generate_new_unique_value()
                # Handle null_count for BlockShuffler (business layer)
                if self.null_count > 0:
                    if self._null_rng is None:
                        self._null_rng = CheckpointablePRNG(seed=self._seed)

                    total_capacity = self.sampler.total_capacity + self.null_count
                    null_probability = self.null_count / total_capacity
                    if self._null_rng.random() < null_probability:
                        return None
                return (
                    int(round(sample_result.value))
                    if sample_result.value is not None
                    else None
                )

            elif phase == SamplingPhase.REPLAY:
                # REPLAY phase: pure new value generation (continues from where PREPARE left off)
                return self._generate_new_unique_value()

        else:
            raise ValueError(f"Unknown constraint_type: {self.constraint_type}")

    def sample(self) -> Optional[int]:
        """
        Sample an integer value (backward compatibility, defaults to PREPARE phase)

        Returns:
            int: Sampled integer, or None for null values
        """
        return self.sample_with_phase(SamplingPhase.PREPARE)

    def convert_scalar_to_business_value(self, shuffler_result) -> int:
        """Convert BlockShuffler result to business value

        Args:
            shuffler_result: The detailed result object from BlockShuffler

        Returns:
            int: The converted business value (integer)
        """
        # For integer types, mainly type conversion
        if shuffler_result is None or shuffler_result.value is None:
            return None
        # Fix floating point precision issue: use truncation instead of rounding
        # to avoid different float values mapping to same integer
        return int(shuffler_result.value)

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

        Args:
            n: Number of samples to generate

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


class FloatSampler:
    """
    Float/Double/Decimal sampler using InverseTransformSampler
    Handles TiDB floating point column statistics
    """

    def __init__(
        self,
        buckets: List[
            TiDBNumericBucket
        ],  # TiDB numeric bucket objects from schema_manager
        top_n: Optional[
            List[TiDBNumericTopNItem]
        ] = None,  # TiDB numeric TopN objects from schema_manager
        null_count: int = 0,
        precision: Optional[int] = None,
        is_cumulative_count: bool = False,  # False for absolute counts
        constraint_type: str = "normal",  # "normal" or "unique"
        seed: Optional[int] = 42,
        column_name: str = None,  # Column name for logging
    ):
        """
        Initialize FloatSampler from standardized TiDB numeric objects

        Args:
            buckets: List of TiDB numeric bucket objects (standardized by schema_manager)
            top_n: List of TiDB numeric TopN objects (standardized by schema_manager)
            null_count: Number of null values
            precision: Number of decimal places (for DECIMAL type)
            is_cumulative_count: False for absolute counts
            constraint_type: "normal" for InverseTransformSampler, "unique" for BlockShuffler
            seed: Random seed
        """
        self.constraint_type = constraint_type
        self.null_count = null_count
        self.precision = precision
        self.column_name = column_name or "unknown_column"
        self._seed = seed
        self._null_rng = CheckpointablePRNG(seed=seed) if null_count > 0 else None

        # Convert TiDB numeric objects to mathematical objects
        math_buckets = [bucket.to_histogram_bucket() for bucket in (buckets or [])]
        math_topn = [item.to_topn_item() for item in (top_n or [])]

        # Choose underlying sampler based on constraint type
        if constraint_type == "unique":
            # Primary key/unique key uses BlockShuffler (no duplicate sampling)
            # Handle case where column has no statistical data
            if not math_buckets and not math_topn:
                # print(f"INFO: Unique float column '{self.column_name}' has no statistical data, using pure new value generation")
                self.sampler = None  # No BlockShuffler needed
            else:
                self.sampler = BlockShuffler(
                    buckets=math_buckets,
                    top_n=math_topn,
                    seed=seed,
                )
            # Initialize new value generation strategy for unique columns
            self._init_new_value_generation(math_buckets, math_topn)
        else:
            # Normal column uses InverseTransformSampler (probability sampling)
            self.sampler = InverseTransformSampler(
                buckets=math_buckets,
                top_n=math_topn,
                null_count=null_count,
                is_cumulative_count=is_cumulative_count,
                seed=seed,
            )

    @property
    def its(self):
        """Backward compatibility: Access underlying InverseTransformSampler"""
        if self.constraint_type == "unique":
            # BlockShuffler mode, return None or raise error for backward compatibility
            raise AttributeError(
                "FloatSampler in 'unique' constraint mode uses BlockShuffler, "
                "not InverseTransformSampler. Use sampler.sampler instead."
            )
        return self.sampler

    def _init_new_value_generation(self, math_buckets, math_topn):
        """Initialize new value generation strategy for unique columns when BlockShuffler is exhausted"""
        # Analyze value range from statistical information
        all_values = []

        # Collect values from TopN
        for topn in math_topn:
            all_values.append(topn.value)

        # Collect values from buckets (use bounds as reference)
        for bucket in math_buckets:
            all_values.extend([bucket.lower_bound, bucket.upper_bound])

        if all_values:
            self._min_value = float(min(all_values))
            self._max_value = float(max(all_values))
            # For floats, increment by minimum precision unit
            if self.precision is not None:
                self._increment = 10 ** (-self.precision)
                self._next_new_value = self._max_value + self._increment
            else:
                self._increment = 1.0
                self._next_new_value = self._max_value + 1.0
        else:
            # No statistical data, start from 1.0
            self._min_value = 1.0
            self._max_value = 1.0
            if self.precision is not None:
                self._increment = 10 ** (-self.precision)
            else:
                self._increment = 1.0
            self._next_new_value = 1.0

        # Flag to log exhaustion message only once
        self._block_exhausted_logged = False

    def _generate_new_unique_value(self) -> float:
        """Generate new unique value using incremental strategy"""
        value = self._next_new_value
        self._next_new_value += self._increment

        # Apply precision if specified
        if self.precision is not None:
            value = round(value, self.precision)

        return value

    def has_statistical_capacity(self) -> bool:
        """Check if statistical information still has capacity"""
        if self.constraint_type == "unique":
            if self.sampler is None:
                return False  # No statistical data, no capacity
            return hasattr(self.sampler, "has_capacity") and self.sampler.has_capacity()
        return True  # Normal columns always have capacity

    def sample_with_phase(self, phase: SamplingPhase) -> Optional[float]:
        """Sample with phase awareness"""
        if self.constraint_type == "normal":
            # Normal columns use InverseTransformSampler in both phases
            sample_result = self.sampler.sample()
            if sample_result is None or sample_result.value is None:
                return None
            # Apply precision if specified
            if self.precision is not None:
                return round(sample_result.value, self.precision)
            return sample_result.value

        elif self.constraint_type == "unique":
            if phase == SamplingPhase.PREPARE:
                # PREPARE phase: BlockShuffler + new value generation
                if self.sampler is None:
                    # No statistical data, use pure new value generation
                    return self._generate_new_unique_value()

                sample_result = self.sampler.sample()
                if sample_result is None:
                    # BlockShuffler exhausted, log once and switch to new value generation
                    if not self._block_exhausted_logged:
                        remaining_capacity = getattr(
                            self.sampler, "remaining_capacity", "unknown"
                        )
                        print(
                            f"INFO: BlockShuffler exhausted for unique float column '{self.column_name}', switching to new value generation. Remaining capacity: {remaining_capacity}"
                        )
                        self._block_exhausted_logged = True
                    return self._generate_new_unique_value()
                # Handle null_count for BlockShuffler (business layer)
                if self.null_count > 0:
                    if self._null_rng is None:
                        self._null_rng = CheckpointablePRNG(seed=self._seed)

                    total_capacity = self.sampler.total_capacity + self.null_count
                    null_probability = self.null_count / total_capacity
                    if self._null_rng.random() < null_probability:
                        return None

                value = sample_result.value
                if value is None:
                    return None
                # Apply precision if specified
                if self.precision is not None:
                    return round(value, self.precision)
                return value

            elif phase == SamplingPhase.REPLAY:
                # REPLAY phase: continue from where PREPARE left off (no duplicates due to state persistence)
                return self._generate_new_unique_value()

        else:
            raise ValueError(f"Unknown constraint_type: {self.constraint_type}")

    def sample(self) -> Optional[float]:
        """
        Sample a float value (backward compatibility, defaults to PREPARE phase)

        Returns:
            float: Sampled float, or None for null values
        """
        return self.sample_with_phase(SamplingPhase.PREPARE)

    def convert_scalar_to_business_value(self, shuffler_result) -> float:
        """Convert BlockShuffler result to business value

        Args:
            shuffler_result: The detailed result object from BlockShuffler

        Returns:
            float: The converted business value (float)
        """
        # For float types, mainly type conversion and precision handling
        if shuffler_result is None or shuffler_result.value is None:
            return None
        value = shuffler_result.value
        # Apply precision if specified
        if self.precision is not None:
            return round(value, self.precision)
        return float(value)

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

        Args:
            n: Number of samples to generate

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


# Alias for backward compatibility
DecimalSampler = FloatSampler