from dataclasses import InitVar, dataclass, field
from typing import Any, Callable, Dict, List

import numpy as np
import pandas as pd
from jinja2 import Template

from samplers.sampling_phase import SamplingPhase
from samplers.state_manager import StateManager

from .column import Column
from .graph import Graph
from .index import Index
from .partition import Partition

DDL_TEMPLATE = """
CREATE TABLE IF NOT EXISTS `{{ schema_name }}`.`{{ table_name }}` (
    {%- for c in columns %}
    `{{c.name}}` {{c.data_type}}
    {%- if c.character_maximum_length %}({{c.character_maximum_length}}){%- endif -%}
    {%- if c.is_nullable == "NO" %} NOT NULL{%- endif -%}
    {%- if not c.column_default %}
        {%- if c.is_nullable == "YES"  %} DEFAULT NULL{%- endif %}
    {%- elif c.data_type in ["char", "varchar", "text", "date", "datetime"] %} DEFAULT "{{c.column_default}}"
    {%- else %} DEFAULT {{c.column_default}}
    {%- endif -%}
    {%- if columns|length != c.ordinal_position -%},{% endif %}
    {%- endfor %}
    {%- if indexes|length > 0 %},{% endif %}
    {% for idx in indexes -%}
    {%- if idx.key_name == "PRIMARY" -%}
    PRIMARY KEY (
        {%- for c in idx.columns -%}
        `{{c.name}}`{% if loop.index != idx.columns|length %},{% endif %}
        {%- endfor -%}
    )
    {%- if idx.clustered == "YES" %} CLUSTERED{% else %} NONCLUSTERED{% endif %}
    {%- else %}
    {% if idx.unique %}UNIQUE {% endif %}KEY `{{idx.key_name}}` (
        {%- for c in idx.columns -%}
        `{{c.name}}`{% if loop.index != idx.columns|length %},{% endif %}
        {%- endfor -%}
    )
    {%- if idx.is_global %} GLOBAL{% endif -%}
    {%- endif -%}
    {%- if loop.index != indexes|length -%},{%- endif -%}
    {%- endfor %}
)
{%- if partitions|length > 0 %}
PARTITION BY {{partitions[0].partition_method}} ({{partitions[0].partition_expression}})
{%- if partitions[0].partition_method in ["KEY", "HASH"] %}
PARTITIONS {{partitions|length}}
{%- else %}
(
    {%- for p in partitions %}
    PARTITION {{p.partition_name}}
    {%- if p.partition_method in ["LIST", "LIST COLUMNS"] %} VALUES IN({{p.partition_description}}){% endif -%}
    {%- if p.partition_method in ["RANGE", "RANGE COLUMNS"] %} VALUES LESS THAN({{p.partition_description}}){% endif -%}
    {%- if partitions|length != p.partition_ordinal_position -%},{% endif %}
    {%- endfor %}
)
{% endif %}
{% endif %}
"""


@dataclass
class Table:
    name: str
    columns: List[Column] = field(init=False)
    indexes: List[Index] = field(init=False)
    composite_indexes: List[Index] = field(init=False)
    partitions: List[Partition] = field(init=False)
    estimated_row_count: int = field(init=False)
    dataframes: InitVar[Dict[str, pd.DataFrame]]
    graph: InitVar[Graph]
    column_factory: InitVar[Callable[..., Column]]
    index_factory: InitVar[Callable[..., Index]]
    partition_factory: InitVar[Callable[..., Partition]]
    reference_map: InitVar[Dict[str, Dict]] = None

    def __post_init__(
        self,
        dataframes,
        graph,
        column_factory,
        index_factory,
        partition_factory,
        reference_map,
    ):
        # Use provided reference_map (built at Database level)
        if reference_map is None:
            reference_map = {}

        self.columns = self._init_columns(
            columns_df=dataframes["columns_df"],
            stats_buckets_df=dataframes["stats_buckets_df"],
            stats_top_n_df=dataframes["stats_top_n_df"],
            graph=graph,
            column_factory=column_factory,
            reference_map=reference_map,
        )
        self.indexes = self._init_indexes(
            indexes_df=dataframes["indexes_df"],
            index_factory=index_factory,
        )
        self.partitions = self._init_partitions(
            partitions_df=dataframes["partitions_df"],
            partition_factory=partition_factory,
        )
        self.composite_indexes = [
            idx for idx in self.indexes if self.is_composite_key(idx)
        ]
        self.estimated_row_count = self._estimate_rows(
            stats_buckets_df=dataframes["stats_buckets_df"],
            stats_top_n_df=dataframes["stats_top_n_df"],
        )

        # Initialize all column samplers (requires indexes to be created first)
        self._initialize_column_samplers()

        # Initialize composite key samplers (requires column samplers to be created first)
        self._initialize_composite_samplers()

    def _init_columns(
        self,
        columns_df,
        stats_buckets_df,
        stats_top_n_df,
        graph,
        column_factory,
        reference_map,
    ) -> List[Column]:
        columns = []
        for c in columns_df.to_dict(orient="records"):
            col = column_factory(
                name=c["column_name"],
                table_schema=c["table_schema"],
                table_name=c["table_name"],
                ordinal_position=c["ordinal_position"],
                data_type=c["data_type"],
                character_maximum_length=(
                    None
                    if not c["character_maximum_length"]
                    or np.isnan(c["character_maximum_length"])
                    else int(c["character_maximum_length"])
                ),
                stats_buckets_df=stats_buckets_df[
                    (stats_buckets_df.table_schema == c["table_schema"])
                    & (stats_buckets_df.table_name == c["table_name"])
                    & (stats_buckets_df.column_name == c["column_name"])
                ],
                stats_top_n_df=stats_top_n_df[
                    (stats_top_n_df.table_schema == c["table_schema"])
                    & (stats_top_n_df.table_name == c["table_name"])
                    & (stats_top_n_df.column_name == c["column_name"])
                ],
                graph=graph,
            )

            # Apply reference context if this column is in the reference map
            column_key = f"{c['table_schema']}.{c['table_name']}.{c['column_name']}"
            if column_key in reference_map:
                ref_info = reference_map[column_key]
                col.set_reference_context(
                    ref_stats_buckets=ref_info["stats_buckets"],
                    ref_stats_top_n=ref_info["stats_top_n"],
                    ref_seed=ref_info["seed"],
                )

            columns.append(col)
        return columns

    def _init_indexes(self, indexes_df, index_factory) -> List[Index]:
        """
        Build indexes from indexes_df with proper column ordering by seq_in_index.
        Should call _init_columns before _init_indexes because index depends on columns.
        """
        if indexes_df.empty:
            return []

        # Sort by key_name and seq_in_index for proper ordering
        sorted_df = indexes_df.sort_values(["key_name", "seq_in_index"])

        key_dict = {}
        for index in sorted_df.to_dict(orient="records"):
            kn = index["key_name"]
            if kn not in key_dict:
                key_dict[kn] = index_factory(
                    key_name=kn,
                    table_schema=index["table_schema"],
                    table_name=index["table_name"],
                    columns=[],
                    is_unique=not index["non_unique"],
                    is_clustered=(index["clustered"] == "YES"),
                    is_global=bool(index["is_global"]),
                )
            key_dict[kn].columns.append(self.lookup_column(index["column_name"]))
        return list(key_dict.values())

    def _init_partitions(self, partitions_df, partition_factory) -> List[Partition]:
        partitions = []
        for d in partitions_df[partitions_df.partition_name.notna()].to_dict(
            orient="records"
        ):
            partitions.append(partition_factory(**d))
        return partitions

    def schema_name(self) -> str:
        """
        The name of the database the table belong to.
        """
        assert len(self.columns) > 0, "table must have at least one column"
        return self.columns[0].table_schema

    def lookup_column(self, name: str) -> Column:
        for c in self.columns:
            if c.name == name:
                return c
        return None

    def sample_row(
        self, phase: SamplingPhase = SamplingPhase.PREPARE
    ) -> Dict[str, Any]:
        """Sample a row with composite key awareness"""
        data = {}
        processed_columns = set()

        # 1. Process unique composite keys (PRIMARY KEY and UNIQUE KEY)
        for composite_idx in self.composite_indexes:
            if composite_idx.is_unique_key:
                # Use phase-aware composite sampling
                composite_values = composite_idx.sample_composite_key_with_phase(phase)
                if composite_values:
                    for i, col in enumerate(composite_idx.columns):
                        data[col.name] = composite_values[i]
                        processed_columns.add(col.name)

        # 2. Process remaining single columns (including columns from non-unique composite indexes)
        for col in self.columns:
            if col.name not in processed_columns:
                data[col.name] = col.sample_for_insert(phase=phase)

        return data

    def show_create_table(self) -> str:
        template = Template(DDL_TEMPLATE)
        return template.render(
            schema_name=self.schema_name(),
            table_name=self.name,
            columns=self.columns,
            indexes=self.indexes,
            partitions=self.partitions,
        )

    def _estimate_rows(self, stats_buckets_df, stats_top_n_df) -> int:
        np_count = (
            stats_buckets_df.groupby("column_name")["count"]
            .sum()
            .add(
                stats_top_n_df.groupby("column_name")["count"].sum(),
                fill_value=0,
            )
            .max()
        )

        return (
            0 if np.isnan(np_count) else int(np_count)
        )  # // 100  # for bank local test

    def _initialize_column_samplers(self):
        """Initialize samplers for each column based on constraints."""
        for col in self.columns:
            if self.is_primary_key(col) or self.is_unique_key(col):
                constraint_type = "unique"  # Use BlockShuffler
            else:
                constraint_type = "normal"  # Use InverseTransformSampler

            col.initialize_sampler(constraint_type)

    def _initialize_composite_samplers(self):
        """Initialize composite key samplers - after column samplers are initialized"""
        for idx in self.composite_indexes:
            if (
                idx.is_unique_key
            ):  # only create composite sampler for unique composite keys
                idx.initialize_composite_sampler()

    # def get_primary_keys(self) -> List[Index]:
    #     """Get the primary key indexes (usually one, but could be multiple)."""
    #     return [index for index in self.indexes if index.key_name == "PRIMARY"]

    # def get_unique_keys(self) -> List[Index]:
    #     """Get all unique key indexes (excluding primary key)."""
    #     return [index for index in self.indexes if index.is_unique and index.key_name != "PRIMARY"]

    # def get_non_unique_indexes(self) -> List[Index]:
    #     """Get all non-unique indexes."""
    #     return [index for index in self.indexes if not index.is_unique]

    def is_composite_key(self, index: Index) -> bool:
        """Check if an index is a composite key (has multiple columns)."""
        return len(index.columns) > 1

    def is_primary_key(self, col: Column) -> bool:
        """Check if a column is part of any primary key."""
        for idx in self.indexes:
            if idx.key_name == "PRIMARY":
                for index_col in idx.columns:
                    if index_col.name == col.name:
                        return True
        return False

    def is_unique_key(self, col: Column) -> bool:
        """Check if a column is part of any unique index (including primary key)."""
        for idx in self.indexes:
            if idx.is_unique:
                for index_col in idx.columns:
                    if index_col.name == col.name:
                        return True
        return False

    def save_all_column_states(self, state_manager: StateManager) -> bool:
        """Save state for all columns in the table"""
        success_count = 0
        total_count = len(self.columns)

        for col in self.columns:
            if col.save_state(state_manager, self.name):
                success_count += 1
            else:
                print(f"⚠️  Failed to save state for column {col.name}")

        if success_count == total_count:
            return True
        else:
            print(
                f"⚠️  Saved {success_count}/{total_count} column states for table {self.name}"
            )
            return False

    def load_all_column_states(self, state_manager: StateManager) -> bool:
        """Load state for all columns in the table"""
        success_count = 0
        found_count = 0

        for col in self.columns:
            if state_manager.has_column_state(self.name, col.name):
                found_count += 1
                if col.load_state(state_manager, self.name):
                    success_count += 1
                else:
                    print(f"⚠️  Failed to load state for column {col.name}")
            else:
                print(f"📄 No saved state found for column {col.name}")

        if found_count == 0:
            print(f"📄 No saved states found for table {self.name}")
            return False
        elif success_count == found_count:
            return True
        else:
            print(
                f"⚠️  Loaded {success_count}/{found_count} available column states for table {self.name}"
            )
            return False

    def has_saved_states(self, state_manager: StateManager) -> bool:
        """Check if the table has any saved column states"""
        for col in self.columns:
            if state_manager.has_column_state(self.name, col.name):
                return True
        return False