"""DSL datatypes for puzzle grammars.

A Grammar is a small, fully-executable description of one puzzle language:
lexicon + affixes + rewrite rules + constituent order (+ numerals). The LLM
proposes grammars as JSON; `from_json` parses defensively (a malformed rule
is dropped, never fatal). `mdl` gives the description length used by the
verifier's simplicity penalty.

JSON shape the proposer LLM emits:
{
  "lexicon":  [{"morph": "kupu", "gloss": "bird", "pos": "N"}, ...],
  "affixes":  [{"position": "prefix"|"suffix"|"circumfix", "form": "na",
                "form2": "", "feature": "DEF", "trigger": "N"}, ...],
  "rewrites": [{"pattern": "a+a", "repl": "aa", "context": ""}, ...]   # regex
  "order":    ["V", "S", "O"],          # target constituent order
  "agree":    [{"src_slot": "S", "dst_slot": "V", "feature": "NUM"}, ...],
  "redup":    [{"scope": "first_syllable", "feature": "PL"}, ...]
}
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional


@dataclass
class LexEntry:
    morph: str
    gloss: str
    pos: str = ""
    features: Dict[str, str] = field(default_factory=dict)


@dataclass
class Affix:
    position: str  # prefix | suffix | infix | circumfix
    form: str
    feature: str = ""  # what it marks, e.g. "PL", "PST", "DEF"
    trigger: str = ""  # pos or feature it attaches to; "" = any
    form2: str = ""  # second part for circumfix


@dataclass
class Rewrite:
    pattern: str  # regex over the surface string
    repl: str
    context: str = ""  # optional regex that must match for rule to fire

    def apply(self, s: str) -> str:
        try:
            if self.context and not re.search(self.context, s):
                return s
            return re.sub(self.pattern, self.repl, s)
        except re.error:
            return s


@dataclass
class Agree:
    src_slot: str
    dst_slot: str
    feature: str


@dataclass
class Redup:
    scope: str = "first_syllable"  # or "full", "first_cv"
    feature: str = "PL"


@dataclass
class Grammar:
    lexicon: List[LexEntry] = field(default_factory=list)
    affixes: List[Affix] = field(default_factory=list)
    rewrites: List[Rewrite] = field(default_factory=list)
    order: List[str] = field(default_factory=list)
    agree: List[Agree] = field(default_factory=list)
    redup: List[Redup] = field(default_factory=list)
    notes: str = ""

    # ---- lookup helpers ----
    def by_gloss(self) -> Dict[str, LexEntry]:
        return {e.gloss.casefold(): e for e in self.lexicon}

    def by_morph(self) -> Dict[str, LexEntry]:
        return {e.morph: e for e in self.lexicon}

    def mdl(self) -> float:
        """Description length: total symbols in the grammar. Lightly weighted
        by the verifier; only breaks ties between equally-fitting grammars."""
        n = 0
        for e in self.lexicon:
            n += len(e.morph) + len(e.gloss) + 2
        for a in self.affixes:
            n += len(a.form) + len(a.form2) + len(a.feature) + 3
        for r in self.rewrites:
            n += len(r.pattern) + len(r.repl) + len(r.context) + 3
        n += 2 * len(self.order) + 4 * len(self.agree) + 4 * len(self.redup)
        return float(n)

    def to_json(self) -> str:
        return json.dumps(
            {
                "lexicon": [
                    {"morph": e.morph, "gloss": e.gloss, "pos": e.pos, "features": e.features}
                    for e in self.lexicon
                ],
                "affixes": [
                    {"position": a.position, "form": a.form, "form2": a.form2,
                     "feature": a.feature, "trigger": a.trigger}
                    for a in self.affixes
                ],
                "rewrites": [
                    {"pattern": r.pattern, "repl": r.repl, "context": r.context}
                    for r in self.rewrites
                ],
                "order": self.order,
                "agree": [
                    {"src_slot": g.src_slot, "dst_slot": g.dst_slot, "feature": g.feature}
                    for g in self.agree
                ],
                "redup": [{"scope": d.scope, "feature": d.feature} for d in self.redup],
            },
            ensure_ascii=False,
        )


def _get(d: dict, *keys: str, default: str = "") -> str:
    for k in keys:
        if k in d and d[k] is not None:
            return str(d[k])
    return default


def from_json(text: str) -> Optional[Grammar]:
    """Parse an LLM-emitted grammar. Tolerates surrounding prose/code fences
    and drops malformed entries instead of failing."""
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if not m:
        return None
    try:
        data = json.loads(m.group(0))
    except json.JSONDecodeError:
        # common LLM damage: trailing commas
        try:
            data = json.loads(re.sub(r",\s*([}\]])", r"\1", m.group(0)))
        except json.JSONDecodeError:
            return None
    if not isinstance(data, dict):
        return None

    g = Grammar()
    for e in data.get("lexicon") or []:
        if isinstance(e, dict):
            morph, gloss = _get(e, "morph", "form", "word"), _get(e, "gloss", "meaning")
            if morph and gloss:
                feats = e.get("features") if isinstance(e.get("features"), dict) else {}
                g.lexicon.append(LexEntry(morph, gloss, _get(e, "pos"), {str(k): str(v) for k, v in (feats or {}).items()}))
    for a in data.get("affixes") or []:
        if isinstance(a, dict):
            form = _get(a, "form")
            pos = _get(a, "position", default="suffix").lower()
            if form and pos in ("prefix", "suffix", "infix", "circumfix"):
                g.affixes.append(Affix(pos, form, _get(a, "feature", "gloss"), _get(a, "trigger"), _get(a, "form2")))
    for r in data.get("rewrites") or []:
        if isinstance(r, dict) and _get(r, "pattern"):
            try:
                re.compile(_get(r, "pattern"))
                if _get(r, "context"):
                    re.compile(_get(r, "context"))
            except re.error:
                continue
            g.rewrites.append(Rewrite(_get(r, "pattern"), _get(r, "repl", "replacement"), _get(r, "context")))
    order = data.get("order") or []
    if isinstance(order, list):
        g.order = [str(x) for x in order]
    for ag in data.get("agree") or []:
        if isinstance(ag, dict) and _get(ag, "feature"):
            g.agree.append(Agree(_get(ag, "src_slot", "src"), _get(ag, "dst_slot", "dst"), _get(ag, "feature")))
    for rd in data.get("redup") or []:
        if isinstance(rd, dict):
            g.redup.append(Redup(_get(rd, "scope", default="first_syllable"), _get(rd, "feature", default="PL")))
    if not g.lexicon and not g.affixes and not g.rewrites:
        return None
    return g