from __future__ import annotations

import re
from dataclasses import dataclass, field
from enum import Enum


class MathStructureKind(str, Enum):
    ASSIGN_SET = "assign_set"
    ASSIGN_NUM = "assign_num"
    SET_LITERAL = "set_literal"
    BARE_NUMBER = "bare_number"


@dataclass(frozen=True)
class MathStructureMatch:
    kind: MathStructureKind
    text: str
    normalized_text: str
    span: tuple[int, int]
    children: dict[str, object] = field(default_factory=dict)
    context_before: str = ""
    context_after: str = ""
    ambiguous: bool = False


_INTEGER_PATTERN = r"[+-]?\d+"

_SPEC: list[tuple[MathStructureKind | str, str]] = [
    (
        MathStructureKind.ASSIGN_SET,
        rf"\b[A-Za-z_]\w*\s*=\s*\{{\s*{_INTEGER_PATTERN}(?:\s*,\s*{_INTEGER_PATTERN})*\s*\}}",
    ),
    (
        MathStructureKind.ASSIGN_NUM,
        rf"\b[A-Za-z_]\w*\s*=\s*{_INTEGER_PATTERN}\b",
    ),
    (
        MathStructureKind.SET_LITERAL,
        rf"\{{\s*{_INTEGER_PATTERN}(?:\s*,\s*{_INTEGER_PATTERN})*\s*\}}",
    ),
    (
        MathStructureKind.BARE_NUMBER,
        rf"(?<![\w.]){_INTEGER_PATTERN}(?![\w.])",
    ),
    ("WORD", r"[A-Za-z_]\w*"),
    ("PUNCT", r"[^\w\s]"),
    ("WS", r"\s+"),
]

_MASTER = re.compile(
    "|".join(
        f"(?P<{kind.value if isinstance(kind, MathStructureKind) else kind}>{pattern})"
        for kind, pattern in _SPEC
    ),
    re.IGNORECASE,
)


def normalize_structure_text(value: str) -> str:
    return re.sub(r"\s+", "", value.strip())


def _extract_integers(value: str) -> list[str]:
    return re.findall(_INTEGER_PATTERN, value)


def _context_window(
    text: str,
    start: int,
    end: int,
    window_chars: int,
) -> tuple[str, str]:
    return (
        text[max(0, start - window_chars) : start].strip(),
        text[end : min(len(text), end + window_chars)].strip(),
    )


def _parse_assign_set(token: str) -> dict[str, object]:
    target, _, value = token.partition("=")
    return {
        "target": target.strip(),
        "value": normalize_structure_text(value),
        "elements": _extract_integers(value),
    }


def _parse_assign_num(token: str) -> dict[str, object]:
    target, _, value = token.partition("=")
    return {
        "target": target.strip(),
        "value": normalize_structure_text(value),
    }


def _parse_set_literal(token: str) -> dict[str, object]:
    return {
        "elements": _extract_integers(token),
    }


def scan_math_structures(
    text: str,
    *,
    window_chars: int = 48,
    max_matches: int = 40,
) -> list[MathStructureMatch]:
    matches: list[MathStructureMatch] = []

    for match in _MASTER.finditer(text):
        group_name = match.lastgroup
        if group_name in {"WORD", "PUNCT", "WS"}:
            continue

        start, end = match.span()
        token = match.group()
        normalized = normalize_structure_text(token)
        context_before, context_after = _context_window(
            text,
            start,
            end,
            window_chars,
        )

        if group_name == MathStructureKind.ASSIGN_SET.value:
            matches.append(
                MathStructureMatch(
                    kind=MathStructureKind.ASSIGN_SET,
                    text=token,
                    normalized_text=normalized,
                    span=(start, end),
                    children=_parse_assign_set(token),
                    context_before=context_before,
                    context_after=context_after,
                    ambiguous=False,
                )
            )
        elif group_name == MathStructureKind.ASSIGN_NUM.value:
            matches.append(
                MathStructureMatch(
                    kind=MathStructureKind.ASSIGN_NUM,
                    text=token,
                    normalized_text=normalized,
                    span=(start, end),
                    children=_parse_assign_num(token),
                    context_before=context_before,
                    context_after=context_after,
                    ambiguous=False,
                )
            )
        elif group_name == MathStructureKind.SET_LITERAL.value:
            matches.append(
                MathStructureMatch(
                    kind=MathStructureKind.SET_LITERAL,
                    text=token,
                    normalized_text=normalized,
                    span=(start, end),
                    children=_parse_set_literal(token),
                    context_before=context_before,
                    context_after=context_after,
                    ambiguous=False,
                )
            )
        elif group_name == MathStructureKind.BARE_NUMBER.value:
            matches.append(
                MathStructureMatch(
                    kind=MathStructureKind.BARE_NUMBER,
                    text=token,
                    normalized_text=normalized,
                    span=(start, end),
                    children={},
                    context_before=context_before,
                    context_after=context_after,
                    ambiguous=True,
                )
            )

        if len(matches) >= max_matches:
            break

    return matches


scan = scan_math_structures