"""Marking — Convert k-OA to SOA over Σ^(k) (Definition 4.4, arXiv 1004.2372)."""

from .soa import SOA
from .grammar import Symbol, Epsilon, Empty, Plus, Star, Optional, Concat, Alt


def mark_koa(G):
    """
    Mark a k-OA G as a SOA over Σ^(k).

    Process nodes in arbitrary order. For the i-th occurrence of label a,
    replace by a^(i) (represented as Symbol('a_i')).

    Returns a SOA H over Σ^(k) such that L(G) = strip(L(H)).
    """
    H = SOA()
    H.src = G.src
    H.sink = G.sink
    H._succ = {n: set(succ) for n, succ in G._succ.items()}
    H._pred = {n: set(pred) for n, pred in G._pred.items()}
    H._label = {}
    H._next = G._next

    counts = {}
    for n in G._succ:
        lab = G._label.get(n)
        if lab is not None and not isinstance(lab, (Empty, Epsilon)) and n not in (G.src, G.sink):
            sym = strip_k(lab)
            key = sym.value if isinstance(sym, Symbol) else str(sym)
            counts[key] = counts.get(key, 0) + 1
            H._label[n] = Symbol(f"{key}_{counts[key]}")
        elif n in (G.src, G.sink):
            H._label[n] = None
        else:
            H._label[n] = lab

    return H


def strip_k(node):
    """Remove k-ORE markers from AST: Symbol('a_1') → Symbol('a'), Symbol('b^(2)') → Symbol('b')."""
    if isinstance(node, Symbol):
        import re
        value = node.value
        value = re.sub(r'_\d+$', '', value)
        value = re.sub(r'\^\(\d+\)$', '', value)
        return Symbol(value)
    if isinstance(node, (Epsilon, Empty)):
        return node
    if isinstance(node, (Plus, Optional, Star)):
        return type(node)(strip_k(node.child))
    if isinstance(node, (Concat, Alt)):
        return type(node)([strip_k(child) for child in node.parts])
    return node