"""Subgraph pattern matching for COGILES colored graphs.

Provides a SMARTS-like pattern language for searching substructures
in colored graphs. Extends COGILES syntax with wildcards and color sets:

- ``*`` matches any color
- ``[RBG]`` matches red, blue, or green
- ``[!RBG]`` matches anything except red, blue, or green

Usage::

    import cogiles

    G = cogiles.parse("R(G)(B)(Y)")
    matches = cogiles.find_pattern(G, "R(*)")
    matches = cogiles.find_pattern(G, "[!R]")

"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import networkx as nx
import parsimonious
from networkx.algorithms import isomorphism

from cogiles.colors import symbol_to_color_info, SYMBOL_MAP, rgb_to_symbol
from cogiles.errors import CogilesParseError
from cogiles.parser import _squeeze_list


# ---------------------------------------------------------------------------
# Pattern grammar — extends COGILES with *, [ABC], [!ABC]
# ---------------------------------------------------------------------------

_pattern_grammar = parsimonious.Grammar(
    r"""
    graph           = (branch / pattern_node / anchor / break)*

    branch          = lpar graph rpar
    lpar            = "("
    rpar            = ")"

    break           = "."
    anchor          = ~r"-[\d]+"
    pattern_node    = wildcard / negated_set / color_set / node
    wildcard        = "*"
    negated_set     = "[!" color_symbol+ "]"
    color_set       = "[" color_symbol+ "]"
    color_symbol    = ~r"(Bl|Br|Gr|Pi|Ol|[RGBYKCMWOPTLI])"
    node            = ~r"(Bl|Br|Gr|Pi|Ol|[RGBYKCMWOPTLI])"
    """
)


# ---------------------------------------------------------------------------
# Match predicates
# ---------------------------------------------------------------------------


def _exact_match(symbol: str) -> Callable[[str], bool]:
    """Predicate: matches exactly one canonical color symbol."""
    target = symbol_to_color_info(symbol).symbol
    return lambda s, _t=target: s == _t


def _any_match() -> Callable[[str], bool]:
    """Predicate: matches any color symbol."""
    return lambda s: True


def _set_match(symbols: set[str]) -> Callable[[str], bool]:
    """Predicate: matches any symbol in the given set."""
    canonical = frozenset(symbol_to_color_info(s).symbol for s in symbols)
    return lambda s, _c=canonical: s in _c


def _negated_set_match(symbols: set[str]) -> Callable[[str], bool]:
    """Predicate: matches any symbol NOT in the given set."""
    canonical = frozenset(symbol_to_color_info(s).symbol for s in symbols)
    return lambda s, _c=canonical: s not in _c


# ---------------------------------------------------------------------------
# Pattern visitor
# ---------------------------------------------------------------------------


class _PatternVisitor(parsimonious.NodeVisitor):
    """Walk the pattern parse tree and build a pattern graph with predicates."""

    def __init__(self):
        super().__init__()
        self.node_index: int = 0
        self.node_predicates: dict[int, Callable] = {}

        self.anchor_map: dict[int, int] = {}
        self.anchor_edges: list[tuple[int, int]] = []

        self.node_sequence_map: dict[int, int] = {}
        self.prev_node: int = 0
        self.is_break: bool = False

    def _register_node(self, predicate: Callable) -> int:
        """Register a new pattern node with a match predicate."""
        idx = self.node_index

        if idx != 0:
            if self.is_break:
                self.is_break = False
            else:
                self.node_sequence_map[idx] = self.prev_node

        self.node_predicates[idx] = predicate
        self.prev_node = idx
        self.node_index += 1
        return idx

    # -- graph / branch / break / anchor (same semantics as CogilesVisitor) --

    def visit_graph(self, node, visited_children):
        return visited_children

    def visit_branch(self, node, visited_children):
        branch = _squeeze_list(visited_children[1:-1])
        if not isinstance(branch, list):
            branch = [branch]

        if not branch:
            return branch

        node_indices = [x for x in branch if isinstance(x, int)]

        if not node_indices:
            raise ValueError("Branch must contain at least one node")

        first_node = node_indices[0]
        if first_node not in self.node_sequence_map:
            raise ValueError("Branch has no parent node to connect to")

        self.prev_node = self.node_sequence_map[first_node]
        self.is_break = False
        return branch

    def visit_break(self, node, visited_children):
        self.is_break = True
        return None

    def visit_anchor(self, node, visited_children):
        anchor_number = int(node.text.replace("-", ""))
        node_index = self.prev_node
        if anchor_number not in self.anchor_map:
            self.anchor_map[anchor_number] = node_index
        else:
            self.anchor_edges.append(
                (self.anchor_map[anchor_number], node_index)
            )
        return None

    # -- pattern-specific node visitors --

    def visit_pattern_node(self, node, visited_children):
        # The matched alternative (wildcard/negated_set/color_set/node)
        # already registered the node via _register_node and returned its idx.
        return visited_children[0]

    def visit_node(self, node, visited_children):
        symbol = str(node.text)
        return self._register_node(_exact_match(symbol))

    def visit_wildcard(self, node, visited_children):
        return self._register_node(_any_match())

    def visit_color_set(self, node, visited_children):
        # children: "[", color_symbol+, "]"
        symbols = visited_children[1]
        if not isinstance(symbols, list):
            symbols = [symbols]
        return self._register_node(_set_match(set(symbols)))

    def visit_negated_set(self, node, visited_children):
        # children: "[!", color_symbol+, "]"
        symbols = visited_children[1]
        if not isinstance(symbols, list):
            symbols = [symbols]
        return self._register_node(_negated_set_match(set(symbols)))

    def visit_color_symbol(self, node, visited_children):
        return str(node.text)

    def generic_visit(self, node, visited_children):
        return visited_children or node

    # -- graph construction --

    def build_graph(self) -> nx.Graph:
        """Construct the pattern graph with ``match`` predicates on nodes."""
        G = nx.Graph()

        for idx in sorted(self.node_predicates):
            G.add_node(idx, match=self.node_predicates[idx])

        for child, parent in self.node_sequence_map.items():
            G.add_edge(child, parent)

        for i, j in self.anchor_edges:
            if i == j:
                raise ValueError(
                    f"Anchor creates a self-loop on node {i}"
                )
            G.add_edge(i, j)

        return G


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class Match:
    """A single pattern match result.

    Attributes
    ----------
    mapping : dict[int, int]
        Pattern node ID -> graph node ID.
    """

    mapping: dict[int, int]

    @property
    def nodes(self) -> frozenset[int]:
        """The set of matched graph node IDs."""
        return frozenset(self.mapping.values())


def parse_pattern(pattern_string: str) -> nx.Graph:
    """Parse a COGILES pattern string into a pattern graph.

    Pattern syntax extends COGILES with:

    - ``*`` — wildcard, matches any color
    - ``[RBG]`` — color set, matches any listed color
    - ``[!RBG]`` — negated set, matches any color NOT listed

    Parameters
    ----------
    pattern_string : str
        The pattern string to parse.

    Returns
    -------
    nx.Graph
        A graph where each node has a ``match`` attribute (a callable
        predicate accepting a color symbol and returning bool).

    Raises
    ------
    CogilesParseError
        If the pattern string is invalid.
    """
    if not isinstance(pattern_string, str):
        raise CogilesParseError(
            f"Expected string, got {type(pattern_string).__name__}",
            text=repr(pattern_string),
        )

    if not pattern_string:
        raise CogilesParseError("Empty pattern string", text=pattern_string)

    try:
        tree = _pattern_grammar.parse(pattern_string)
    except parsimonious.exceptions.ParseError as exc:
        raise CogilesParseError(
            f"Invalid pattern syntax: {exc}",
            text=pattern_string,
            position=getattr(exc, "pos", None),
        ) from exc

    try:
        visitor = _PatternVisitor()
        visitor.visit(tree)
    except parsimonious.exceptions.VisitationError as exc:
        cause = exc.__context__ or exc.__cause__
        msg = str(cause) if cause else str(exc).split("\n")[0]
        raise CogilesParseError(msg, text=pattern_string) from exc

    try:
        return visitor.build_graph()
    except ValueError as exc:
        raise CogilesParseError(str(exc), text=pattern_string) from exc


def _symbol_from_attrs(attrs: dict) -> str | None:
    """Extract canonical color symbol from graph node attributes."""
    symbol = attrs.get("symbol")
    if symbol is not None and symbol in SYMBOL_MAP:
        return SYMBOL_MAP[symbol].symbol
    color = attrs.get("color")
    if color is not None:
        return rgb_to_symbol(tuple(color))
    return None


def find_pattern(
    graph: nx.Graph | str,
    pattern: nx.Graph | str,
) -> list[Match]:
    """Find all occurrences of a pattern in a graph.

    Uses VF2 subgraph isomorphism with color-aware node matching.
    Matches are deduplicated by the set of matched graph nodes.

    Parameters
    ----------
    graph : nx.Graph or str
        The target graph (or COGILES string) to search in.
    pattern : nx.Graph or str
        The pattern graph (or pattern string) to search for.
        Supports wildcards (``*``), color sets (``[RBG]``),
        and negated sets (``[!RBG]``).

    Returns
    -------
    list[Match]
        All distinct matches found. Each :class:`Match` contains
        a ``mapping`` (pattern node -> graph node) and a ``nodes``
        property (set of matched graph node IDs).
    """
    if isinstance(graph, str):
        from cogiles.parser import parse

        graph = parse(graph)

    if isinstance(pattern, str):
        pattern = parse_pattern(pattern)

    if pattern.number_of_nodes() == 0:
        return []

    if pattern.number_of_nodes() > graph.number_of_nodes():
        return []

    def node_match(graph_attrs, pattern_attrs):
        predicate = pattern_attrs.get("match")
        if predicate is None:
            return True
        symbol = _symbol_from_attrs(graph_attrs)
        if symbol is None:
            return False
        return predicate(symbol)

    GM = isomorphism.GraphMatcher(graph, pattern, node_match=node_match)

    matches = []
    seen: set[frozenset[int]] = set()
    for mapping in GM.subgraph_isomorphisms_iter():
        # mapping: graph_node -> pattern_node
        matched_nodes = frozenset(mapping.keys())
        if matched_nodes in seen:
            continue
        seen.add(matched_nodes)
        # Invert to pattern_node -> graph_node
        inv = {v: k for k, v in mapping.items()}
        matches.append(Match(mapping=inv))

    return matches
