"""
State management for sampler persistence across prepare/replay phases
"""

import pickle
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Set

import boto3


@dataclass
class BlockShufflerState:
    """State for BlockShuffler - bottom layer mathematical sampler"""

    # PRNG state (core)
    prng_position: int

    # Capacity and exhaustion state
    remaining_capacity: int
    is_exhausted: bool

    # Block internal states (critical)
    block_states: Dict[int, Dict[str, Any]]  # block_id -> state

    # Sampling sequence state (critical)
    block_id_sequence: List[int]  # shuffled block sampling order
    current_block_index: int  # current sampling position


@dataclass
class BusinessSamplerState:
    """State for business layer sampler"""

    # Bottom layer BlockShuffler state
    block_shuffler_state: Optional[BlockShufflerState]

    # Business layer specific state
    constraint_type: str  # "normal" | "unique"
    column_name: str  # for logging

    # Unique constraint related state
    block_exhausted_logged: bool
    generated_values: Set[Any]  # generated new values for deduplication
    new_value_counter: int  # new value generation counter

    # Null handling state
    null_count: int
    null_probability: float


@dataclass
class ColumnState:
    """Complete column-level state"""

    # Business sampler state
    sampler_state: BusinessSamplerState

    # Column cache state (critical)
    cache_contents: List[Any]  # deque contents
    cache_maxlen: int  # deque maxlen parameter

    # Column metadata
    data_type: str
    character_maximum_length: int
    is_initialized: bool  # whether sampler is initialized


class StateManager:
    """Handles state persistence for columns"""

    def __init__(self, state_config):
        # Handle both dict and object config formats
        if isinstance(state_config, dict):
            backend_type = state_config.get("backend", "local")
        else:
            backend_type = getattr(state_config, "backend", "local")

        if backend_type == "s3":
            self.backend_type = "s3"

            # Handle both dict and object config formats for S3 settings
            if isinstance(state_config, dict):
                s3_config = state_config.get("s3", {})
                self.bucket = s3_config.get("bucket")
                base_dir = s3_config.get("base_dir")
                if base_dir:
                    self.save_path = f"{base_dir.rstrip('/')}/states"
                else:
                    self.save_path = state_config.get("save_path", "states")
            else:
                self.bucket = getattr(state_config.s3, "bucket", None)
                base_dir = getattr(state_config.s3, "base_dir", None)
                if base_dir:
                    self.save_path = f"{base_dir.rstrip('/')}/states"
                else:
                    self.save_path = getattr(state_config, "save_path", "states")

            self.s3_client = boto3.client("s3")
        else:
            self.backend_type = "local"

            # Handle both dict and object config formats for save_path
            if isinstance(state_config, dict):
                save_path_str = state_config.get("save_path", "./states")
            else:
                save_path_str = getattr(state_config, "save_path", "./states")

            self.save_path = Path(save_path_str)
            self.save_path.mkdir(parents=True, exist_ok=True)

    def get_column_state_file(self, table_name: str, column_name: str) -> str:
        """Get the state file path/key for a specific column"""
        if self.backend_type == "s3":
            # For S3, use clean key without ./ prefix
            save_path = self.save_path.lstrip("./")
            return f"{save_path}/{table_name}_{column_name}_state.pkl"
        else:
            # For local files, keep the original path
            return f"{self.save_path}/{table_name}_{column_name}_state.pkl"

    def save_column_state(
        self, table_name: str, column_name: str, state: ColumnState
    ) -> bool:
        """Save column state to file or S3"""
        try:
            pickled_data = pickle.dumps(state)

            if self.backend_type == "s3":
                state_key = self.get_column_state_file(table_name, column_name)
                self.s3_client.put_object(
                    Bucket=self.bucket, Key=state_key, Body=pickled_data
                )
            else:
                state_file = self.get_column_state_file(table_name, column_name)
                with open(state_file, "wb") as f:
                    f.write(pickled_data)
            return True
        except Exception as e:
            print(f"❌ Failed to save state for {table_name}.{column_name}: {e}")
            return False

    def load_column_state(
        self, table_name: str, column_name: str
    ) -> Optional[ColumnState]:
        """Load column state from file or S3"""
        try:
            if self.backend_type == "s3":
                state_key = self.get_column_state_file(table_name, column_name)
                response = self.s3_client.get_object(Bucket=self.bucket, Key=state_key)
                pickled_data = response["Body"].read()
            else:
                state_file = Path(self.get_column_state_file(table_name, column_name))
                if not state_file.exists():
                    return None
                with open(state_file, "rb") as f:
                    pickled_data = f.read()

            return pickle.loads(pickled_data)
        except Exception as e:
            print(f"❌ Failed to load state for {table_name}.{column_name}: {e}")
            return None

    def has_column_state(self, table_name: str, column_name: str) -> bool:
        """Check if column state file/object exists"""
        try:
            if self.backend_type == "s3":
                state_key = self.get_column_state_file(table_name, column_name)
                self.s3_client.head_object(Bucket=self.bucket, Key=state_key)
                return True
            else:
                return Path(
                    self.get_column_state_file(table_name, column_name)
                ).exists()
        except Exception as e:
            print(f"❌ Failed to check state for {table_name}.{column_name}: {e}")
            return False