from __future__ import annotations

from dataclasses import dataclass
from enum import IntEnum
from typing import TYPE_CHECKING

from blends.stack.attributes import (
    STACK_GRAPH_IMPORT_MODULE_STEM,
    STACK_GRAPH_IMPORT_SOURCE_SYMBOL,
    STACK_GRAPH_IS_DEFINITION,
    STACK_GRAPH_IS_EXPORTED,
    STACK_GRAPH_IS_IMPORT_BINDING,
    STACK_GRAPH_IS_REFERENCE,
    STACK_GRAPH_KIND,
    STACK_GRAPH_PRECEDENCE,
    STACK_GRAPH_SCOPE,
    STACK_GRAPH_SCOPE_TYPE,
    STACK_GRAPH_SYMBOL,
)
from blends.stack.node_kinds import (
    StackGraphNodeKind,
)

if TYPE_CHECKING:
    from blends.models import (
        NId,
        SyntaxGraph,
    )

ROOT_INDEX: int = 0
JUMP_TO_INDEX: int = 1


class Degree(IntEnum):
    ZERO = 0
    ONE = 1
    MULTIPLE = 2


def _nid_sort_key(n_id: NId) -> tuple[int, int | str]:
    if n_id.isdecimal():
        return (0, int(n_id))
    return (1, n_id)


def _resolve_file_path(graph: SyntaxGraph, path: str | None) -> str:
    if path is not None:
        return path
    metadata = graph.nodes.get("0", {})
    node_path = metadata.get("path")
    return node_path if isinstance(node_path, str) else ""


def _collect_stack_graph_node_ids(graph: SyntaxGraph) -> list[NId]:
    return [
        node_id for node_id in graph.nodes if graph.nodes[node_id].get(STACK_GRAPH_KIND) is not None
    ]


@dataclass(slots=True)
class NodeSemantics:
    node_kind: list[str | None]
    node_symbol_id: list[int | None]
    node_scope_index: list[int | None]
    node_precedence: list[int]
    exported_scopes: set[int]
    node_is_reference: list[bool]
    node_is_definition: list[bool]
    node_scope_type: list[str | None]
    node_is_import_binding: list[bool]
    node_import_module_stem: list[str | None]
    node_import_source_symbol: list[str | None]


def _find_root_nid(graph: SyntaxGraph, stack_node_ids: list[NId]) -> NId | None:
    for nid in stack_node_ids:
        if graph.nodes[nid].get(STACK_GRAPH_KIND) == StackGraphNodeKind.ROOT.value:
            return nid
    return None


def _build_node_index_maps(
    stack_node_ids: list[NId], *, root_nid: NId | None
) -> tuple[dict[NId, int], list[NId]]:
    nid_to_index: dict[NId, int] = {}
    index_to_nid: list[NId] = []

    if root_nid is not None:
        nid_to_index[root_nid] = ROOT_INDEX
    index_to_nid.append(root_nid or "")

    index_to_nid.append("")

    for node_id in stack_node_ids:
        if node_id in nid_to_index:
            continue
        nid_to_index[node_id] = len(index_to_nid)
        index_to_nid.append(node_id)
    return nid_to_index, index_to_nid


def _fill_structural_attrs(
    attrs: dict[str, object],
    idx: int,
    nid_to_index: dict[NId, int],
    symbols: SymbolInterner,
    sem: NodeSemantics,
) -> None:
    kind = attrs.get(STACK_GRAPH_KIND)
    if isinstance(kind, str):
        sem.node_kind[idx] = kind

    symbol = attrs.get(STACK_GRAPH_SYMBOL)
    if isinstance(symbol, str) and symbol:
        sem.node_symbol_id[idx] = symbols.intern(symbol)

    scope = attrs.get(STACK_GRAPH_SCOPE)
    if isinstance(scope, str):
        sem.node_scope_index[idx] = nid_to_index.get(scope)

    precedence = attrs.get(STACK_GRAPH_PRECEDENCE, 0)
    sem.node_precedence[idx] = precedence if isinstance(precedence, int) else 0


def _fill_flag_attrs(
    attrs: dict[str, object],
    idx: int,
    sem: NodeSemantics,
) -> None:
    if attrs.get(STACK_GRAPH_IS_EXPORTED) is True:
        sem.exported_scopes.add(idx)

    is_reference = attrs.get(STACK_GRAPH_IS_REFERENCE)
    if isinstance(is_reference, bool):
        sem.node_is_reference[idx] = is_reference

    is_definition = attrs.get(STACK_GRAPH_IS_DEFINITION)
    if isinstance(is_definition, bool):
        sem.node_is_definition[idx] = is_definition

    scope_type = attrs.get(STACK_GRAPH_SCOPE_TYPE)
    if isinstance(scope_type, str):
        sem.node_scope_type[idx] = scope_type

    is_import_binding = attrs.get(STACK_GRAPH_IS_IMPORT_BINDING)
    if isinstance(is_import_binding, bool):
        sem.node_is_import_binding[idx] = is_import_binding

    import_module_stem = attrs.get(STACK_GRAPH_IMPORT_MODULE_STEM)
    if isinstance(import_module_stem, str):
        sem.node_import_module_stem[idx] = import_module_stem

    import_source_symbol = attrs.get(STACK_GRAPH_IMPORT_SOURCE_SYMBOL)
    if isinstance(import_source_symbol, str):
        sem.node_import_source_symbol[idx] = import_source_symbol


def _extract_node_semantics(
    graph: SyntaxGraph,
    *,
    stack_node_ids: list[NId],
    nid_to_index: dict[NId, int],
    symbols: SymbolInterner,
    node_count: int,
) -> NodeSemantics:
    sem = NodeSemantics(
        node_kind=[None] * node_count,
        node_symbol_id=[None] * node_count,
        node_scope_index=[None] * node_count,
        node_precedence=[0] * node_count,
        exported_scopes=set(),
        node_is_reference=[False] * node_count,
        node_is_definition=[False] * node_count,
        node_scope_type=[None] * node_count,
        node_is_import_binding=[False] * node_count,
        node_import_module_stem=[None] * node_count,
        node_import_source_symbol=[None] * node_count,
    )
    sem.node_kind[ROOT_INDEX] = StackGraphNodeKind.ROOT.value
    sem.node_kind[JUMP_TO_INDEX] = StackGraphNodeKind.JUMP_TO.value

    for node_id in stack_node_ids:
        attrs = graph.nodes[node_id]
        idx = nid_to_index[node_id]
        _fill_structural_attrs(attrs, idx, nid_to_index, symbols, sem)
        _fill_flag_attrs(attrs, idx, sem)

    return sem


def _extract_and_normalize_outgoing(
    graph: SyntaxGraph,
    *,
    stack_node_ids: list[NId],
    nid_to_index: dict[NId, int],
    node_count: int,
) -> tuple[list[list[tuple[int, int]]], list[Degree]]:
    outgoing: list[list[tuple[int, int]]] = [[] for _ in range(node_count)]
    incoming_counts: list[int] = [0] * node_count
    stack_node_id_set = set(stack_node_ids)
    for source_id, sink_id, edge_attrs in graph.edges(data=True):
        if source_id not in stack_node_id_set or sink_id not in stack_node_id_set:
            continue
        if "precedence" not in edge_attrs:
            continue
        source_index = nid_to_index[source_id]
        sink_index = nid_to_index[sink_id]
        precedence = edge_attrs.get("precedence", 0)
        outgoing[source_index].append(
            (
                sink_index,
                precedence if isinstance(precedence, int) else 0,
            )
        )
        incoming_counts[sink_index] += 1

    normalized_outgoing: list[list[tuple[int, int]]] = [[] for _ in range(node_count)]
    for source_index, edges in enumerate(outgoing):
        if not edges:
            continue
        edges.sort(key=lambda item: (item[0], item[1]))
        normalized_outgoing[source_index] = edges

    incoming_degree = [
        Degree.ZERO if count == 0 else Degree.ONE if count == 1 else Degree.MULTIPLE
        for count in incoming_counts
    ]
    return normalized_outgoing, incoming_degree


def _build_file_maps(
    file_path: str, *, stack_node_ids: list[NId], nid_to_index: dict[NId, int], node_count: int
) -> tuple[dict[str, frozenset[int]], list[str]]:
    real_node_indices = frozenset(nid_to_index[node_id] for node_id in stack_node_ids)
    file_to_nodes = {file_path: real_node_indices}
    node_to_file = [file_path if index in real_node_indices else "" for index in range(node_count)]
    return file_to_nodes, node_to_file


@dataclass(slots=True)
class SymbolInterner:
    symbol_to_id: dict[str, int]
    id_to_symbol: list[str]

    def intern(self, symbol: str) -> int:
        symbol_id = self.symbol_to_id.get(symbol)
        if symbol_id is not None:
            return symbol_id
        symbol_id = len(self.id_to_symbol)
        self.symbol_to_id[symbol] = symbol_id
        self.id_to_symbol.append(symbol)
        return symbol_id

    def lookup(self, symbol_id: int) -> str:
        return self.id_to_symbol[symbol_id]


@dataclass(frozen=True, slots=True)
class StackGraphView:
    file_path: str
    nid_to_index: dict[NId, int]
    index_to_nid: list[NId]
    symbols: SymbolInterner
    node_kind: list[str | None]
    node_symbol_id: list[int | None]
    node_scope_index: list[int | None]
    node_precedence: list[int]
    outgoing: list[list[tuple[int, int]]]
    incoming_degree: list[Degree]
    exported_scopes: set[int]
    file_to_nodes: dict[str, frozenset[int]]
    node_to_file: list[str]
    node_is_reference: list[bool]
    node_is_definition: list[bool]
    node_scope_type: list[str | None]
    node_is_import_binding: list[bool]
    node_import_module_stem: list[str | None]
    node_import_source_symbol: list[str | None]

    @classmethod
    def from_syntax_graph(cls, graph: SyntaxGraph, *, path: str | None = None) -> StackGraphView:
        file_path = _resolve_file_path(graph, path)

        stack_node_ids = sorted(_collect_stack_graph_node_ids(graph), key=_nid_sort_key)
        root_nid = _find_root_nid(graph, stack_node_ids)
        nid_to_index, index_to_nid = _build_node_index_maps(stack_node_ids, root_nid=root_nid)

        symbols = SymbolInterner(symbol_to_id={}, id_to_symbol=[])

        node_count = len(index_to_nid)
        semantics = _extract_node_semantics(
            graph,
            stack_node_ids=stack_node_ids,
            nid_to_index=nid_to_index,
            symbols=symbols,
            node_count=node_count,
        )
        outgoing, incoming_degree = _extract_and_normalize_outgoing(
            graph,
            stack_node_ids=stack_node_ids,
            nid_to_index=nid_to_index,
            node_count=node_count,
        )
        file_to_nodes, node_to_file = _build_file_maps(
            file_path,
            stack_node_ids=stack_node_ids,
            nid_to_index=nid_to_index,
            node_count=node_count,
        )

        return cls(
            file_path=file_path,
            nid_to_index=nid_to_index,
            index_to_nid=index_to_nid,
            symbols=symbols,
            node_kind=semantics.node_kind,
            node_symbol_id=semantics.node_symbol_id,
            node_scope_index=semantics.node_scope_index,
            node_precedence=semantics.node_precedence,
            outgoing=outgoing,
            incoming_degree=incoming_degree,
            exported_scopes=semantics.exported_scopes,
            file_to_nodes=file_to_nodes,
            node_to_file=node_to_file,
            node_is_reference=semantics.node_is_reference,
            node_is_definition=semantics.node_is_definition,
            node_scope_type=semantics.node_scope_type,
            node_is_import_binding=semantics.node_is_import_binding,
            node_import_module_stem=semantics.node_import_module_stem,
            node_import_source_symbol=semantics.node_import_source_symbol,
        )

    def kind_at(self, node_index: int) -> str:
        if node_index == 0:
            return StackGraphNodeKind.ROOT.value
        if node_index == 1:
            return StackGraphNodeKind.JUMP_TO.value
        kind = self.node_kind[node_index]
        return kind if isinstance(kind, str) else ""

    def symbol_id_at(self, node_index: int) -> int | None:
        if node_index in {0, 1}:
            return None
        return self.node_symbol_id[node_index]

    def scope_index_at(self, node_index: int) -> int | None:
        if node_index in {0, 1}:
            return None
        return self.node_scope_index[node_index]
