from __future__ import annotations

from collections import Counter, defaultdict
from dataclasses import dataclass, field

from ...metamodel.metamodel import Aggregate, Exists, Logical, Match, Not, Task, Union, Update
from .graph import SQL2GraphArtifacts, SQL2GraphNode
from .table_shape import SQL2TableShapeSidecar


@dataclass(frozen=True)
class SQL2WriteChannel:
    channel_id: int
    node_id: int
    target_relation_id: int
    target_relation_name: str
    output_columns: tuple[str, ...]
    key_var_ids: tuple[int, ...]
    rule: Task


@dataclass(frozen=True)
class SQL2CoalescedGroup:
    group_id: str
    stratum_id: int
    table_name: str
    key_var_ids: tuple[int, ...]
    output_columns: tuple[str, ...]
    key_columns: tuple[str, ...]
    node_ids: tuple[int, ...]
    channels: tuple[SQL2WriteChannel, ...]
    shared_tasks: tuple[Task, ...] = field(default_factory=tuple)
    coalesced_rule: Task | None = None


@dataclass(frozen=True)
class SQL2CoalescedArtifacts:
    groups: tuple[SQL2CoalescedGroup, ...]
    groups_by_stratum: dict[int, tuple[SQL2CoalescedGroup, ...]]
    group_for_node_id: dict[int, str]
    program_by_stratum: dict[int, Task]
    program_root: Task | None


class SQL2Coalescer:
    """Coalesce graph nodes by `(stratum, table, key signature)` with channel isolation."""

    def coalesce(
        self,
        *,
        graph: SQL2GraphArtifacts,
        table_shape: SQL2TableShapeSidecar,
    ) -> SQL2CoalescedArtifacts:
        nodes_by_id = {node.node_id: node for node in graph.nodes}
        grouped_node_ids: dict[tuple[int, str, tuple[int, ...]], list[int]] = defaultdict(list)
        for node in graph.nodes:
            stratum = graph.stratum_by_node.get(node.node_id, 0)
            table_name = table_shape.table_name_by_relation_id.get(node.target_relation_id, node.target_table_name)
            grouped_node_ids[(stratum, table_name, node.key_var_ids)].append(node.node_id)

        groups: list[SQL2CoalescedGroup] = []
        groups_by_stratum: dict[int, list[SQL2CoalescedGroup]] = defaultdict(list)
        group_for_node_id: dict[int, str] = {}

        for (stratum_id, table_name, key_var_ids) in sorted(
            grouped_node_ids,
            key=lambda k: (k[0], k[1], k[2]),
        ):
            node_ids = sorted(grouped_node_ids[(stratum_id, table_name, key_var_ids)])
            channels: list[SQL2WriteChannel] = []
            for node_id in node_ids:
                node = nodes_by_id[node_id]
                if node.rewritten_rule is None:
                    continue
                channels.append(
                    SQL2WriteChannel(
                        channel_id=node_id,
                        node_id=node.node_id,
                        target_relation_id=node.target_relation_id,
                        target_relation_name=node.target_relation_name,
                        output_columns=self._output_columns(node=node, table_shape=table_shape),
                        key_var_ids=node.key_var_ids,
                        rule=node.rewritten_rule,
                    )
                )

            shared_tasks = self._shared_tasks(channels)
            coalesced_rule = (
                Logical(body=tuple(channel.rule for channel in channels), scope=True)
                if channels
                else None
            )
            output_columns, key_columns = self._group_output_and_key_columns(
                channels=channels,
                table_shape=table_shape,
            )
            group_id = self._group_id(stratum_id=stratum_id, table_name=table_name, key_var_ids=key_var_ids)
            group = SQL2CoalescedGroup(
                group_id=group_id,
                stratum_id=stratum_id,
                table_name=table_name,
                key_var_ids=key_var_ids,
                output_columns=output_columns,
                key_columns=key_columns,
                node_ids=tuple(node_ids),
                channels=tuple(channels),
                shared_tasks=shared_tasks,
                coalesced_rule=coalesced_rule,
            )
            groups.append(group)
            groups_by_stratum[stratum_id].append(group)
            for node_id in node_ids:
                group_for_node_id[node_id] = group_id

        ordered_groups = tuple(groups)
        group_by_id = {group.group_id: group for group in ordered_groups}
        ordered_groups_by_stratum = {
            stratum_id: self._order_groups_for_stratum(
                stratum_id=stratum_id,
                groups=tuple(stratum_groups),
                graph=graph,
                group_for_node_id=group_for_node_id,
                group_by_id=group_by_id,
            )
            for stratum_id, stratum_groups in sorted(groups_by_stratum.items())
        }
        program_by_stratum: dict[int, Task] = {}
        for stratum_id, stratum_groups in ordered_groups_by_stratum.items():
            rules = tuple(
                group.coalesced_rule
                for group in stratum_groups
                if group.coalesced_rule is not None
            )
            program_by_stratum[stratum_id] = Logical(body=rules, scope=True)

        program_root = (
            Logical(
                body=tuple(program_by_stratum[stratum_id] for stratum_id in sorted(program_by_stratum)),
                scope=True,
            )
            if program_by_stratum
            else None
        )
        return SQL2CoalescedArtifacts(
            groups=ordered_groups,
            groups_by_stratum=ordered_groups_by_stratum,
            group_for_node_id=group_for_node_id,
            program_by_stratum=program_by_stratum,
            program_root=program_root,
        )

    def _order_groups_for_stratum(
        self,
        *,
        stratum_id: int,
        groups: tuple[SQL2CoalescedGroup, ...],
        graph: SQL2GraphArtifacts,
        group_for_node_id: dict[int, str],
        group_by_id: dict[str, SQL2CoalescedGroup],
    ) -> tuple[SQL2CoalescedGroup, ...]:
        group_ids = {group.group_id for group in groups}
        indegree: dict[str, int] = {group.group_id: 0 for group in groups}
        outgoing: dict[str, set[str]] = {group.group_id: set() for group in groups}

        for edge in graph.edges:
            source_group_id = group_for_node_id.get(edge.source_node_id)
            target_group_id = group_for_node_id.get(edge.target_node_id)
            if source_group_id is None or target_group_id is None:
                continue
            if source_group_id == target_group_id:
                continue
            if source_group_id not in group_ids or target_group_id not in group_ids:
                continue
            source_group = group_by_id[source_group_id]
            target_group = group_by_id[target_group_id]
            if source_group.stratum_id != stratum_id or target_group.stratum_id != stratum_id:
                continue
            if target_group_id in outgoing[source_group_id]:
                continue
            outgoing[source_group_id].add(target_group_id)
            indegree[target_group_id] += 1

        def sort_key(group_id: str) -> tuple[int, str, tuple[int, ...], str]:
            group = group_by_id[group_id]
            return (
                group.node_ids[0] if group.node_ids else -1,
                group.table_name,
                group.key_var_ids,
                group.group_id,
            )

        ready = sorted((group_id for group_id, deg in indegree.items() if deg == 0), key=sort_key)
        ordered_group_ids: list[str] = []
        while ready:
            current = ready.pop(0)
            ordered_group_ids.append(current)
            for nxt in sorted(outgoing[current], key=sort_key):
                indegree[nxt] -= 1
                if indegree[nxt] == 0:
                    ready.append(nxt)
                    ready.sort(key=sort_key)

        if len(ordered_group_ids) != len(groups):
            # Fall back deterministically if dependency graph is unexpectedly cyclic.
            ordered_group_ids = sorted(group_ids, key=sort_key)
        return tuple(group_by_id[group_id] for group_id in ordered_group_ids)

    def _group_id(self, *, stratum_id: int, table_name: str, key_var_ids: tuple[int, ...]) -> str:
        key_part = ",".join(str(var_id) for var_id in key_var_ids)
        return f"s{stratum_id}:{table_name}:{key_part}"

    def _group_output_and_key_columns(
        self,
        *,
        channels: list[SQL2WriteChannel],
        table_shape: SQL2TableShapeSidecar,
    ) -> tuple[tuple[str, ...], tuple[str, ...]]:
        if not channels:
            return (), ()
        table = table_shape.table_for_relation_id(channels[0].target_relation_id)
        if table is not None:
            return (
                tuple(column.column_name for column in table.columns),
                tuple(table.key_column_names),
            )
        output_columns = self._dedupe_columns(
            tuple(column_name for channel in channels for column_name in channel.output_columns)
        )
        return output_columns, ()

    def _dedupe_columns(self, columns: tuple[str, ...]) -> tuple[str, ...]:
        return tuple(dict.fromkeys(columns))

    def _output_columns(
        self,
        *,
        node: SQL2GraphNode,
        table_shape: SQL2TableShapeSidecar,
    ) -> tuple[str, ...]:
        mapped_column = table_shape.property_column_by_relation_id.get(node.target_relation_id)
        if mapped_column is not None:
            return (mapped_column,)
        table = table_shape.table_for_relation_id(node.target_relation_id)
        if table is None:
            return self._output_columns_from_rule(node)
        if node.target_relation_id == table.owner_relation_id:
            return tuple(column.column_name for column in table.columns)
        return tuple(
            column.column_name
            for column in table.columns
            if column.source_relation_id == node.target_relation_id
        )

    def _output_columns_from_rule(self, node: SQL2GraphNode) -> tuple[str, ...]:
        update = self._find_update_for_relation(node.rewritten_rule, node.target_relation_id)
        if update is None:
            return ()
        return tuple(field.name for field in update.relation.fields)

    def _find_update_for_relation(self, task: Task | None, relation_id: int) -> Update | None:
        if task is None:
            return None
        if isinstance(task, Update) and task.relation.id == relation_id:
            return task
        if isinstance(task, Logical):
            for child in task.body:
                update = self._find_update_for_relation(child, relation_id)
                if update is not None:
                    return update
            return None
        if isinstance(task, Not):
            return self._find_update_for_relation(task.task, relation_id)
        if isinstance(task, Exists):
            return self._find_update_for_relation(task.task, relation_id)
        if isinstance(task, Aggregate):
            return self._find_update_for_relation(task.body, relation_id)
        if isinstance(task, Union):
            for child in task.tasks:
                update = self._find_update_for_relation(child, relation_id)
                if update is not None:
                    return update
            return None
        if isinstance(task, Match):
            for child in task.tasks:
                update = self._find_update_for_relation(child, relation_id)
                if update is not None:
                    return update
        return None

    def _shared_tasks(self, channels: list[SQL2WriteChannel]) -> tuple[Task, ...]:
        counts: Counter[Task] = Counter()
        for channel in channels:
            for task in self._rule_tasks(channel.rule):
                counts[task] += 1
        shared = [task for task, count in counts.items() if count > 1]
        shared.sort(key=lambda task: task.id)
        return tuple(shared)

    def _rule_tasks(self, rule: Task) -> tuple[Task, ...]:
        if isinstance(rule, Logical):
            return tuple(rule.body)
        return (rule,)
