"""Blocksworld: generator, verifier, prompts.

Pure planning domain — no code syntax, no arithmetic. Plans are symbolically
verifiable by simulation, so STaR bucketing works. Instances are generated
(contamination-free by construction) with difficulty = blocks + moves.
"""

import json
import random
import re

MOVE_RE = re.compile(
    r"move\s+([A-Z])\s+(?:onto|on top of|on)\s+(?:the\s+)?(table|[A-Z])",
    re.IGNORECASE)


def gen_instance(n_blocks, rng):
    blocks = [chr(65 + i) for i in range(n_blocks)]

    def random_stacks():
        bs = blocks[:]
        rng.shuffle(bs)
        stacks, i = [], 0
        while i < len(bs):
            take = rng.randint(1, len(bs) - i)
            stacks.append(bs[i : i + take])
            i += take
        return stacks

    init = random_stacks()
    goal = random_stacks()
    while goal == init:
        goal = random_stacks()
    return {"blocks": blocks, "init": init, "goal": goal}


def fmt_state(stacks):
    out = []
    for st in stacks:
        if len(st) == 1:
            out.append(f"{st[0]} is on the table")
        else:
            out.append(f"{st[0]} is on the table with "
                       + " on top, then ".join(
                           [f"{b}" for b in st[1:]]) + " on top")
    # clearer explicit form
    lines = []
    for st in stacks:
        lines.append(f"stack: {' -> '.join(st)} (bottom -> top)")
    return "; ".join(lines)


def question(inst):
    return (
        "You are stacking blocks. Only the TOP block of a stack can be "
        "moved, one block at a time.\n"
        f"Blocks: {', '.join(inst['blocks'])}\n"
        f"Initial state: {fmt_state(inst['init'])}\n"
        f"Goal state: {fmt_state(inst['goal'])}\n"
        "Give a plan as a numbered list of moves, each exactly of the form "
        "'move X onto Y' or 'move X onto the table'.")


def verify_plan(inst, text, max_moves=40):
    stacks = [st[:] for st in inst["init"]]

    def top_of(b):
        for st in stacks:
            if st and st[-1] == b:
                return st
        return None

    moves = MOVE_RE.findall(text)
    if not moves or len(moves) > max_moves:
        return False
    for b, tgt in moves:
        b = b.upper()
        tgt = tgt if tgt.lower() == "table" else tgt.upper()
        src = top_of(b)
        if src is None:
            return False  # b not clear (or nonexistent)
        if tgt == "table" or tgt.lower() == "table":
            src.pop()
            stacks.append([b])
        else:
            dst = top_of(tgt)
            if dst is None or b == tgt:
                return False
            src.pop()
            dst.append(b)
        stacks = [st for st in stacks if st]
    norm = sorted(tuple(st) for st in stacks)
    return norm == sorted(tuple(st) for st in inst["goal"])


def make_dataset(n_train=400, n_test=200, seed=0):
    rng = random.Random(seed)
    items = []
    for split, n in (("train", n_train), ("test", n_test)):
        for i in range(n):
            nb = rng.choice([3, 3, 4, 4, 5])
            inst = gen_instance(nb, rng)
            items.append({"split": split, "task_id": f"bw_{split}_{i}",
                          "n_blocks": nb, **inst,
                          "question": question(inst)})
    return items


if __name__ == "__main__":
    ds = make_dataset()
    print(json.dumps(ds[0], indent=1))
    # verifier self-test: identity plan on trivial instance
    inst = {"blocks": ["A", "B"], "init": [["A"], ["B"]],
            "goal": [["A", "B"]]}
    assert verify_plan(inst, "1. move B onto A")
    assert not verify_plan(inst, "1. move A onto A")
    print("verifier self-test ok")