"""Validation parser for positional-format GSM8K examples.

Implements the six checks from Check-in 2's Appendix (Step 4):
  (a) required fields present in every node
      - entry node:  STATE, GOAL, PREDICT
      - other nodes: MOVE, OBSERVE, STATE, PREDICT
  (b) node ids sequential starting from 0
  (c) entry node contains a GOAL field
  (d) final node's PREDICT is "goal resolved"
  (e) <answer> matches the original GSM8K #### answer
  (f) look-ahead detection: numbers in a PREDICT field must already appear
      in the question or in a prior (or same) node's OBSERVE/STATE fields

Usage:
  from validator import validate
  ok, errors = validate(positional_text, original_answer="72")
"""
import re

NODE_RE = re.compile(r'<node\s+id="(\d+)"\s+type="(\w+)">(.*?)</node>', re.S)
ANSWER_RE = re.compile(r"<answer>(.*?)</answer>", re.S)
FIELD_NAMES = ["STATE", "GOAL", "MOVE", "OBSERVE", "PREDICT"]
NUM_RE = re.compile(r"\d+(?:\.\d+)?")
VALID_TYPES = {"entry", "filter", "resolve", "combine", "check", "branch"}


def _fields(body):
    """Split a node body into {FIELD: text} using field-name anchors."""
    positions = []
    for f in FIELD_NAMES:
        m = re.search(rf"^\s*{f}:", body, re.M)
        if m:
            positions.append((m.start(), f))
    positions.sort()
    out = {}
    for (start, f), nxt in zip(positions, positions[1:] + [(len(body), None)]):
        text = body[start:nxt[0]]
        out[f] = text.split(":", 1)[1].strip()
    return out


def _nums(text):
    return set(NUM_RE.findall(text or ""))


def validate(text, original_answer, question=""):
    errors = []
    nodes = NODE_RE.findall(text)
    if not nodes:
        return False, ["no <node> blocks found"]

    # (b) sequential ids from 0
    ids = [int(i) for i, _, _ in nodes]
    if ids != list(range(len(ids))):
        errors.append(f"(b) node ids not sequential from 0: {ids}")

    parsed = []
    for nid, ntype, body in nodes:
        if ntype not in VALID_TYPES:
            errors.append(f"(a) node {nid}: unknown type '{ntype}'")
        parsed.append((int(nid), ntype, _fields(body)))

    # (a) required fields per node
    for nid, ntype, f in parsed:
        req = {"STATE", "GOAL", "PREDICT"} if ntype == "entry" else {"MOVE", "OBSERVE", "STATE", "PREDICT"}
        missing = req - set(f)
        if missing:
            errors.append(f"(a) node {nid} ({ntype}): missing {sorted(missing)}")

    # (c) entry node has GOAL
    if not any(ntype == "entry" and "GOAL" in f for _, ntype, f in parsed):
        errors.append("(c) no entry node with GOAL field")

    # (d) final node PREDICT == goal resolved
    if parsed:
        last = parsed[-1][2].get("PREDICT", "")
        if "goal resolved" not in last.lower():
            errors.append(f"(d) final PREDICT is not 'goal resolved': {last!r}")

    # (e) answer matches
    m = ANSWER_RE.search(text)
    if not m:
        errors.append("(e) no <answer> tag")
    else:
        got = m.group(1).strip().replace(",", "").replace("$", "")
        want = str(original_answer).strip().replace(",", "").replace("$", "")
        if got != want:
            errors.append(f"(e) answer mismatch: got {got!r}, want {want!r}")

    # (f) look-ahead detection
    seen = _nums(question)
    for nid, ntype, f in parsed:
        seen |= _nums(f.get("STATE", "")) | _nums(f.get("OBSERVE", ""))
        lookahead = _nums(f.get("PREDICT", "")) - seen
        if lookahead:
            errors.append(f"(f) node {nid}: PREDICT contains unseen numbers {sorted(lookahead)}")

    return (not errors), errors


if __name__ == "__main__":
    import json, sys
    data = json.load(open(sys.argv[1]))
    passed = failed = 0
    for ex in data:
        ok, errs = validate(ex["positional"], ex["answer"], ex.get("question", ""))
        if ok:
            passed += 1
        else:
            failed += 1
            print(f"FAIL [{ex.get('id','?')}]: {errs}")
    print(f"\n{passed} passed, {failed} failed")