"""Encode NetworkX graphs into canonical COGILES strings."""

from __future__ import annotations

import warnings

import networkx as nx

from cogiles.colors import rgb_to_symbol, SYMBOL_MAP
from cogiles.errors import CogilesEncodeError


class CogilesEncoder:
    """Encode a NetworkX graph into a canonical COGILES string.

    Algorithm:
    1. Build a mutable adjacency dict from the graph.
    2. Greedy DFS starting from the lowest unvisited node index.
    3. When a node has >1 unanchored neighbors, pre-assign an anchor.
    4. Multiple children are wrapped in parentheses as branches.
    5. Disconnected components are joined with ".".
    """

    def __init__(self, graph: nx.Graph, *, strict: bool = False):
        self.graph = graph
        self.strict = strict

        self.nodes = sorted(graph.nodes())

        # Build mutable adjacency (skip self-loops)
        self.adjacency: dict[int, set[int]] = {n: set() for n in self.nodes}
        for u, v in graph.edges():
            if u != v:
                self.adjacency[u].add(v)
                self.adjacency[v].add(u)

        self.node_visited: dict[int, bool] = {n: False for n in self.nodes}
        self.current_anchor_index: int = 1
        self.index_anchor_map: dict[int, int] = {}

    def encode(self) -> str:
        """Encode the graph and return the COGILES string."""
        components: list[str] = []

        index = self._get_unvisited()
        while index is not None:
            components.append(self._encode_branch(index))
            index = self._get_unvisited()

        return ".".join(components)

    def _get_unvisited(self) -> int | None:
        """Return the lowest unvisited node index, or None."""
        for n in self.nodes:
            if not self.node_visited[n]:
                return n
        return None

    def _encode_branch(self, node_index: int) -> str:
        """Recursively encode the connected component from node_index."""
        # Bug fix (E-BUG-4): check symbol BEFORE color so that a node with
        # a valid symbol but color=None can still be encoded.
        symbol = self.graph.nodes[node_index].get("symbol")
        if symbol is not None and symbol in SYMBOL_MAP:
            # Bug fix (E-BUG-3): canonicalize known aliases.
            # Example: symbol='Bl' is encoded as 'K' (the canonical form
            # for black) so that encode output is always idempotent.
            value = SYMBOL_MAP[symbol].symbol
        else:
            color = self.graph.nodes[node_index].get("color")
            if color is None:
                raise CogilesEncodeError(
                    f"Node {node_index} has no 'color' attribute",
                    node=node_index,
                )
            # Bug fix (E-BUG-1/2): if symbol exists but is not in the
            # color registry (e.g., 'XX', ''), warn and fall back to
            # color-based matching instead of blindly trusting it.
            if symbol is not None:
                warnings.warn(
                    f"Node {node_index} has unrecognized symbol "
                    f"{symbol!r}, falling back to color matching",
                    stacklevel=2,
                )
            value = rgb_to_symbol(
                tuple(color), strict=self.strict, node=node_index
            )

        neighbors = sorted(self.adjacency[node_index])
        neighbors_anchor = [n for n in neighbors if n in self.index_anchor_map]
        neighbors_normal = [n for n in neighbors if n not in self.index_anchor_map]

        if len(neighbors_normal) > 1:
            value += self._create_anchor(node_index)

        for neighbor in neighbors_anchor + neighbors_normal:
            self.adjacency[node_index].discard(neighbor)
            self.adjacency[neighbor].discard(node_index)

            if self.node_visited[neighbor]:
                continue

            if neighbor in self.index_anchor_map:
                value += f"-{self.index_anchor_map[neighbor]}"
                continue

            if len(neighbors) == 1:
                value += self._encode_branch(neighbor)
            else:
                value += f"({self._encode_branch(neighbor)})"

        self.node_visited[node_index] = True
        return value

    def _create_anchor(self, node_index: int) -> str:
        """Create and register a new anchor for the given node."""
        anchor_num = self.current_anchor_index
        self.index_anchor_map[node_index] = anchor_num
        self.current_anchor_index += 1
        return f"-{anchor_num}"


def encode(graph: nx.Graph, *, strict: bool = False) -> str:
    """Encode a NetworkX graph into a canonical COGILES string.

    Parameters
    ----------
    graph : nx.Graph
        A graph where each node has a ``color`` attribute (RGB tuple).
        Optionally, nodes may have a ``symbol`` attribute used directly.
    strict : bool, optional
        If True, raise ``CogilesEncodeError`` when a node's color does
        not exactly match any known color. Default: nearest match.

    Returns
    -------
    str
        The canonical COGILES string.

    Raises
    ------
    CogilesEncodeError
        If ``strict=True`` and a node has an unrecognized color, or
        if a node is missing the ``color`` attribute.
    """
    if graph.number_of_nodes() == 0:
        return ""

    encoder = CogilesEncoder(graph, strict=strict)
    return encoder.encode()
