from __future__ import annotations

from collections import defaultdict

from auraindex.models import ColumnUsage, ParsedQueryFeatures, WorkloadStats
from auraindex.utils.text import dedupe_preserve_order


def build_workload_stats(parsed_queries: list[ParsedQueryFeatures]) -> WorkloadStats:
    stats = WorkloadStats(total_queries=len(parsed_queries))
    read_kinds = {"select"}
    write_kinds = {"insert", "update", "delete", "merge", "replace"}

    for query in parsed_queries:
        if query.statement_kind not in read_kinds.union(write_kinds):
            continue
        query_tables = set(query.tables) if query.tables else {"_unknown"}
        for table in query_tables:
            stats.table_query_count[table] = stats.table_query_count.get(table, 0) + 1
            if query.statement_kind in read_kinds:
                stats.table_read_count[table] = stats.table_read_count.get(table, 0) + 1
            elif query.statement_kind in write_kinds:
                stats.table_write_count[table] = stats.table_write_count.get(table, 0) + 1

        _increment_role_counts(stats, query.where_columns, role="where")
        _increment_role_counts(stats, query.join_columns, role="join")
        _increment_role_counts(stats, query.order_by_columns, role="order")
        _increment_role_counts(stats, query.group_by_columns, role="group")

        for table, columns in query.equality_columns_by_table.items():
            ordered = dedupe_preserve_order(columns)
            if len(ordered) < 2:
                continue
            key = tuple(ordered[:3])
            table_map = stats.composite_filters.setdefault(table, {})
            table_map[key] = table_map.get(key, 0) + 1

        where_by_table = _group_columns_by_table(query.where_columns)
        order_by_table = _group_columns_by_table(query.order_by_columns)
        for table in set(where_by_table).intersection(order_by_table):
            where_cols = tuple(dedupe_preserve_order(where_by_table[table]))
            order_cols = tuple(dedupe_preserve_order(order_by_table[table]))
            if not where_cols or not order_cols:
                continue
            pattern_key = (where_cols, order_cols)
            table_patterns = stats.where_order_patterns.setdefault(table, {})
            table_patterns[pattern_key] = table_patterns.get(pattern_key, 0) + 1

    return stats


def _increment_role_counts(stats: WorkloadStats, pairs: list[tuple[str, str]], role: str) -> None:
    for table, column in set(pairs):
        table_usage = stats.column_usage.setdefault(table, {})
        usage = table_usage.setdefault(column, ColumnUsage())

        if role == "where":
            usage.where_count += 1
        elif role == "join":
            usage.join_count += 1
        elif role == "order":
            usage.order_count += 1
        elif role == "group":
            usage.group_count += 1


def _group_columns_by_table(pairs: list[tuple[str, str]]) -> dict[str, list[str]]:
    grouped: dict[str, list[str]] = defaultdict(list)
    for table, column in pairs:
        grouped[table].append(column)
    return grouped
