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

import pandas as pd

from .graph import Graph
from .table import Table


@dataclass
class Database:
    name: str
    tables: List[Table] = field(init=False)
    dataframes: InitVar[Dict[str, pd.DataFrame]]
    graph: InitVar[Graph]
    table_factory: InitVar[Callable[..., Table]]

    def __post_init__(self, dataframes, graph, table_factory):
        self.tables = []
        columns_df = dataframes["columns_df"]
        partitions_df = dataframes["partitions_df"]
        indexes_df = dataframes["indexes_df"]
        stats_buckets_df = dataframes["stats_buckets_df"]
        stats_top_n_df = dataframes["stats_top_n_df"]

        # Build global reference map using all tables' data
        reference_map = self._build_global_reference_map(
            graph=graph,
            indexes_df=indexes_df,
            stats_buckets_df=stats_buckets_df,
            stats_top_n_df=stats_top_n_df,
        )

        for table_name in columns_df.table_name.unique():
            self.tables.append(
                table_factory(
                    name=table_name,
                    dataframes={
                        "columns_df": columns_df[columns_df.table_name == table_name],
                        "partitions_df": partitions_df[
                            partitions_df.table_name == table_name
                        ],
                        "indexes_df": indexes_df[indexes_df.table_name == table_name],
                        "stats_buckets_df": stats_buckets_df[
                            stats_buckets_df.table_name == table_name
                        ],
                        "stats_top_n_df": stats_top_n_df[
                            stats_top_n_df.table_name == table_name
                        ],
                    },
                    graph=graph,
                    reference_map=reference_map,
                )
            )

    def _build_global_reference_map(
        self, graph, indexes_df, stats_buckets_df, stats_top_n_df
    ) -> Dict[str, Dict]:
        """Build global reference map for all correlated columns across tables

        For each correlation cluster, select a primary reference column and
        map other columns to use its statistics and seed.

        Selection priority:
        1. Single-column unique key
        2. Composite unique key (prefer table with highest estimated_row_count)
        3. No sharing if no unique key exists

        Returns:
            Dict mapping column_key -> {stats_buckets, stats_top_n, seed}
        """
        reference_map = {}

        # Iterate through all clusters
        for cluster in graph.iter_clusters():
            # Skip single-column clusters
            if len(cluster) <= 1:
                continue

            # Select primary reference column for this cluster (across all tables)
            primary_col = self._select_primary_reference(
                cluster, indexes_df, stats_buckets_df, stats_top_n_df
            )

            if primary_col is None:
                # No suitable primary found, skip sharing
                continue

            # Derive seed from cluster identity
            cluster_hash = graph.get_column_identity(primary_col)
            seed = int(cluster_hash[:8], 16) % (2**31)

            # Get primary column's statistics
            primary_parts = primary_col.split(".")
            if len(primary_parts) != 3:
                continue
            primary_schema, primary_table, primary_col_name = primary_parts

            primary_buckets = stats_buckets_df[
                (stats_buckets_df.table_schema == primary_schema)
                & (stats_buckets_df.table_name == primary_table)
                & (stats_buckets_df.column_name == primary_col_name)
            ]
            primary_topn = stats_top_n_df[
                (stats_top_n_df.table_schema == primary_schema)
                & (stats_top_n_df.table_name == primary_table)
                & (stats_top_n_df.column_name == primary_col_name)
            ]

            # Map all columns in cluster to reference info
            for col_key in cluster:
                if col_key == primary_col:
                    # Primary column: only set seed, use its own stats
                    reference_map[col_key] = {
                        "stats_buckets": None,  # Use own stats
                        "stats_top_n": None,  # Use own stats
                        "seed": seed,
                    }
                else:
                    # Reference column: use primary's stats and seed
                    reference_map[col_key] = {
                        "stats_buckets": primary_buckets,
                        "stats_top_n": primary_topn,
                        "seed": seed,
                    }

        return reference_map

    def _select_primary_reference(
        self, cluster, indexes_df, stats_buckets_df, stats_top_n_df
    ):
        """Select primary reference column from correlation cluster

        Priority:
        1. Single-column unique key
        2. Composite unique key from table with highest row count
        3. None if no unique key
        """
        if indexes_df.empty:
            return None

        # Collect candidate columns with their index info
        candidates = []

        for col_key in cluster:
            parts = col_key.split(".")
            if len(parts) != 3:
                continue
            schema, table, column = parts

            # Find indexes containing this column
            col_indexes = indexes_df[
                (indexes_df.table_schema == schema)
                & (indexes_df.table_name == table)
                & (indexes_df.column_name == column)
            ]

            for _, idx_row in col_indexes.iterrows():
                is_unique = (idx_row["key_name"] == "PRIMARY") or (
                    idx_row["non_unique"] == 0
                )
                if not is_unique:
                    continue

                # Count columns in this index
                index_cols = indexes_df[
                    (indexes_df.table_schema == schema)
                    & (indexes_df.table_name == table)
                    & (indexes_df.key_name == idx_row["key_name"])
                ]
                is_composite = len(index_cols) > 1

                # Estimate row count for this table
                table_buckets = stats_buckets_df[
                    (stats_buckets_df.table_schema == schema)
                    & (stats_buckets_df.table_name == table)
                ]
                table_topn = stats_top_n_df[
                    (stats_top_n_df.table_schema == schema)
                    & (stats_top_n_df.table_name == table)
                ]
                row_count = (
                    table_buckets.groupby("column_name")["count"]
                    .sum()
                    .add(
                        table_topn.groupby("column_name")["count"].sum(),
                        fill_value=0,
                    )
                    .max()
                )
                row_count = 0 if pd.isna(row_count) else int(row_count)

                candidates.append(
                    {
                        "col_key": col_key,
                        "is_composite": is_composite,
                        "row_count": row_count,
                    }
                )

        if not candidates:
            return None

        # Sort by priority: single-column unique first, then by row count
        candidates.sort(key=lambda x: (x["is_composite"], -x["row_count"]))

        return candidates[0]["col_key"]

    def get_table(self, name: str) -> Table:
        for t in self.tables:
            if t.name == name:
                return t
        return None