from typing import List, Optional

from .block_shuffler import BlockShuffler, CompositeBlockShuffler
from .sampling_phase import SamplingPhase


class CompositeBusinessShuffler:
    """Business layer composite shuffler - coordinates business logic of each column"""

    def __init__(self, columns: List):
        """
        Initialize CompositeBusinessShuffler

        Args:
            columns: List of Column objects with initialized samplers
        """
        self.columns = columns
        self.column_samplers = []

        # Collect business samplers from each column
        for col in columns:
            # print(f"DEBUG: CompositeBusinessShuffler collecting column {col.table_name}.{col.name}")
            if hasattr(col, "sampler") and col.sampler is not None:
                self.column_samplers.append(col.sampler)
            else:
                raise ValueError(f"Column {col.name} has no business sampler")

        # Create math layer CompositeBlockShuffler from each column's BlockShuffler
        block_shufflers = []
        for col_sampler in self.column_samplers:
            if hasattr(col_sampler, "sampler") and isinstance(
                col_sampler.sampler, BlockShuffler
            ):
                block_shufflers.append(col_sampler.sampler)
            else:
                raise ValueError(
                    f"Column sampler for {getattr(col_sampler, 'column_name', 'unknown')} "
                    f"doesn't have BlockShuffler. Got type: {type(col_sampler.sampler) if hasattr(col_sampler, 'sampler') else 'no sampler'}"
                )

        self.math_shuffler = CompositeBlockShuffler(shufflers=block_shufflers)

    def sample(self) -> Optional[tuple]:
        """Sample business values (backward compatibility - defaults to PREPARE phase)"""
        return self.sample_with_phase(SamplingPhase.PREPARE)

    def sample_with_phase(self, phase: SamplingPhase) -> Optional[tuple]:
        """Sample business values with phase awareness"""
        if phase == SamplingPhase.PREPARE:
            # PREPARE phase: Use CompositeBlockShuffler for joint sampling
            return self._sample_from_composite_shuffler()

        elif phase == SamplingPhase.REPLAY:
            # REPLAY phase: Each column independently generates new values
            return self._sample_independent_columns(phase)

        else:
            raise ValueError(f"Unknown sampling phase: {phase}")

    def _sample_from_composite_shuffler(self) -> Optional[tuple]:
        """Sample from CompositeBlockShuffler (PREPARE phase logic)"""
        # 1. Get scalar values and detailed results from math layer
        math_result = self.math_shuffler.sample()
        if math_result is None:
            # print(f"DEBUG: CompositeBlockShuffler exhausted! Remaining: {self.math_shuffler.remaining_capacity}")
            return None

        # 2. Let each column's business sampler handle scalar → business value conversion
        business_values = []
        for i, col_sampler in enumerate(self.column_samplers):
            # Get the detailed result for this column
            shuffler_result = (
                math_result.shuffler_results[i]
                if hasattr(math_result, "shuffler_results")
                and math_result.shuffler_results
                else None
            )
            business_value = col_sampler.convert_scalar_to_business_value(
                shuffler_result
            )

            # If any column returns None, the entire composite key is invalid
            # This handles the case where individual column's BlockShuffler is exhausted
            if business_value is None:
                return None

            business_values.append(business_value)

        final_tuple = tuple(business_values)

        return final_tuple

    def _sample_independent_columns(self, phase: SamplingPhase) -> Optional[tuple]:
        """Sample from each column independently (REPLAY phase logic)"""
        business_values = []

        for col_sampler in self.column_samplers:
            # Each column generates new values independently
            if hasattr(col_sampler, "sample_with_phase"):
                business_value = col_sampler.sample_with_phase(phase)
            else:
                # Fallback for samplers without phase awareness
                business_value = col_sampler.sample()

            business_values.append(business_value)

        return tuple(business_values)

    def has_capacity(self) -> bool:
        """Check if the composite shuffler has remaining capacity"""
        return self.math_shuffler.has_capacity()

    def reset(self):
        """Reset the composite shuffler"""
        return self.math_shuffler.reset()

    @property
    def remaining_capacity(self):
        """Get remaining capacity from the underlying math shuffler"""
        return self.math_shuffler.remaining_capacity

    @property
    def total_capacity(self):
        """Get total capacity from the underlying math shuffler"""
        return self.math_shuffler.total_capacity