cleveragents-core/features/steps/plan_explain_steps.py

"""Step definitions for plan_explain.feature.

Tests the plan explain and plan tree CLI formatting and tree-building
logic directly, without requiring a database or full CLI process.
"""

from __future__ import annotations

import json

from behave import given, then, when  # type: ignore[import-untyped]
from behave.runner import Context
from ulid import ULID

from cleveragents.cli.commands.plan import (
    _build_explain_dict,
    _get_decision_label,
    build_decision_tree,
)
from cleveragents.cli.formatting import format_output
from cleveragents.domain.models.core.decision import (
    ContextSnapshot,
    Decision,
    DecisionType,
    ResourceRef,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_PLAN_ID = str(ULID())


def _make_decision(**overrides: object) -> Decision:
    """Build a Decision with sensible defaults."""
    defaults: dict[str, object] = {
        "plan_id": _PLAN_ID,
        "sequence_number": 0,
        "decision_type": DecisionType.PROMPT_DEFINITION,
        "question": "What should we build?",
        "chosen_option": "A REST API",
    }
    dt = overrides.get("decision_type", defaults["decision_type"])
    if dt != DecisionType.PROMPT_DEFINITION and "parent_decision_id" not in overrides:
        defaults["parent_decision_id"] = str(ULID())
    defaults.update(overrides)
    return Decision(**defaults)


# ---------------------------------------------------------------------------
# Given steps - explain
# ---------------------------------------------------------------------------


@given("a test decision for explain")
def step_test_decision(context: Context) -> None:
    context.pe_decision = _make_decision()


@given("a test decision with context snapshot for explain")
def step_test_decision_with_context(context: Context) -> None:
    snap = ContextSnapshot(
        hot_context_hash="sha256:testctx",
        hot_context_ref="store://test",
        relevant_resources=[
            ResourceRef(resource_id=str(ULID()), path="src/main.py"),
        ],
        actor_state_ref="checkpoint://test/001",
    )
    context.pe_decision = _make_decision(context_snapshot=snap)


@given("a test decision with reasoning for explain")
def step_test_decision_with_reasoning(context: Context) -> None:
    context.pe_decision = _make_decision(
        rationale="Chose REST API for simplicity",
        actor_reasoning="The LLM considered multiple approaches...",
    )


@given("a test decision with alternatives for explain")
def step_test_decision_with_alternatives(context: Context) -> None:
    context.pe_decision = _make_decision(
        chosen_option="GraphQL API",
        alternatives_considered=["GraphQL API", "gRPC service"],
    )


@given("a non-existent decision id")
def step_nonexistent_decision(context: Context) -> None:
    context.pe_missing_id = str(ULID())
    # Also store a known list so we can verify it's absent
    context.pe_known_ids = [str(ULID()) for _ in range(3)]


# ---------------------------------------------------------------------------
# Given steps - tree
# ---------------------------------------------------------------------------


@given("a set of test decisions forming a tree")
def step_tree_decisions(context: Context) -> None:
    root_id = str(ULID())
    child1_id = str(ULID())
    child2_id = str(ULID())
    context.pe_decisions = [
        Decision(
            decision_id=root_id,
            plan_id=_PLAN_ID,
            sequence_number=0,
            decision_type=DecisionType.PROMPT_DEFINITION,
            question="What to build?",
            chosen_option="REST API",
        ),
        Decision(
            decision_id=child1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=1,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question="Which framework?",
            chosen_option="FastAPI",
        ),
        Decision(
            decision_id=child2_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=2,
            decision_type=DecisionType.RESOURCE_SELECTION,
            question="Which database?",
            chosen_option="PostgreSQL",
        ),
    ]


@given("a set of test decisions with superseded entries")
def step_tree_with_superseded(context: Context) -> None:
    root_id = str(ULID())
    child_id = str(ULID())
    superseding_id = str(ULID())
    context.pe_decisions = [
        Decision(
            decision_id=root_id,
            plan_id=_PLAN_ID,
            sequence_number=0,
            decision_type=DecisionType.PROMPT_DEFINITION,
            question="What to build?",
            chosen_option="REST API",
        ),
        Decision(
            decision_id=child_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=1,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question="Which framework?",
            chosen_option="Flask",
            superseded_by=superseding_id,
        ),
        Decision(
            decision_id=superseding_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=2,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question="Which framework?",
            chosen_option="FastAPI",
            is_correction=True,
            corrects_decision_id=child_id,
            correction_reason="FastAPI is faster",
        ),
    ]
    context.pe_superseded_id = child_id


@given("a set of test decisions forming a deep tree")
def step_deep_tree(context: Context) -> None:
    root_id = str(ULID())
    child_id = str(ULID())
    grandchild_id = str(ULID())
    context.pe_decisions = [
        Decision(
            decision_id=root_id,
            plan_id=_PLAN_ID,
            sequence_number=0,
            decision_type=DecisionType.PROMPT_DEFINITION,
            question="Root question?",
            chosen_option="Root answer",
        ),
        Decision(
            decision_id=child_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=1,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question="Child question?",
            chosen_option="Child answer",
        ),
        Decision(
            decision_id=grandchild_id,
            plan_id=_PLAN_ID,
            parent_decision_id=child_id,
            sequence_number=2,
            decision_type=DecisionType.IMPLEMENTATION_CHOICE,
            question="Grandchild question?",
            chosen_option="Grandchild answer",
        ),
    ]


@given("an empty list of decisions")
def step_empty_decisions(context: Context) -> None:
    context.pe_decisions = []


# ---------------------------------------------------------------------------
# When steps - explain
# ---------------------------------------------------------------------------


@when("I build the explain dict with default options")
def step_build_explain_default(context: Context) -> None:
    context.pe_explain_dict = _build_explain_dict(context.pe_decision)


@when("I build the explain dict with show-context enabled")
def step_build_explain_context(context: Context) -> None:
    context.pe_explain_dict = _build_explain_dict(
        context.pe_decision, show_context=True
    )


@when("I build the explain dict with show-reasoning enabled")
def step_build_explain_reasoning(context: Context) -> None:
    context.pe_explain_dict = _build_explain_dict(
        context.pe_decision, show_reasoning=True
    )


def _capture_format_output(data, fmt):
    """Call format_output capturing stdout (machine-readable formats write there)."""
    from contextlib import redirect_stdout
    from io import StringIO

    buf = StringIO()
    with redirect_stdout(buf):
        result = format_output(data, fmt)
    return result or buf.getvalue().rstrip("\n")


@when("I format the explain dict as json")
def step_format_explain_json(context: Context) -> None:
    data = _build_explain_dict(context.pe_decision)
    context.pe_json_output = _capture_format_output(data, "json")


@when("I format the explain dict as yaml")
def step_format_explain_yaml(context: Context) -> None:
    data = _build_explain_dict(context.pe_decision)
    context.pe_yaml_output = _capture_format_output(data, "yaml")


# ---------------------------------------------------------------------------
# When steps - tree
# ---------------------------------------------------------------------------


@when("I build the decision tree with default options")
def step_build_tree_default(context: Context) -> None:
    context.pe_tree = build_decision_tree(context.pe_decisions)


@when("I build the decision tree with show-superseded enabled")
def step_build_tree_superseded(context: Context) -> None:
    context.pe_tree = build_decision_tree(context.pe_decisions, show_superseded=True)


@when("I build the decision tree with depth {depth_val:d}")
def step_build_tree_depth(context: Context, depth_val: int) -> None:
    context.pe_tree = build_decision_tree(context.pe_decisions, max_depth=depth_val)


@when("I build the decision tree with default options from empty list")
def step_build_tree_empty(context: Context) -> None:
    context.pe_tree = build_decision_tree(context.pe_decisions)


@when("I format the tree as json")
def step_format_tree_json(context: Context) -> None:
    tree = build_decision_tree(context.pe_decisions)
    context.pe_tree_json = _capture_format_output(tree, "json")


@when("I format the tree as yaml")
def step_format_tree_yaml(context: Context) -> None:
    tree = build_decision_tree(context.pe_decisions)
    context.pe_tree_yaml = _capture_format_output(tree, "yaml")


# ---------------------------------------------------------------------------
# Then steps - explain
# ---------------------------------------------------------------------------


@then('the explain dict should contain key "{key}"')
def step_explain_has_key(context: Context, key: str) -> None:
    assert key in context.pe_explain_dict, (
        f"Expected key '{key}' in {list(context.pe_explain_dict.keys())}"
    )


@then('the explain dict should not contain key "{key}"')
def step_explain_not_has_key(context: Context, key: str) -> None:
    assert key not in context.pe_explain_dict, f"Key '{key}' should not be present"


@then('the context snapshot should contain key "{key}"')
def step_snapshot_has_key(context: Context, key: str) -> None:
    snap = context.pe_explain_dict["context_snapshot"]
    assert key in snap, f"Expected key '{key}' in snapshot: {list(snap.keys())}"


@then("the alternatives list should have {count:d} items")
def step_alternatives_count(context: Context, count: int) -> None:
    alts = context.pe_explain_dict["alternatives"]
    assert len(alts) == count, f"Expected {count} alternatives, got {len(alts)}"


@then('each alternative should have keys "index", "description", and "chosen"')
def step_alternatives_have_required_keys(context: Context) -> None:
    alts = context.pe_explain_dict["alternatives"]
    for i, alt in enumerate(alts):
        assert isinstance(alt, dict), f"Alternative {i} is not a dict: {alt!r}"
        assert "index" in alt, (
            f"Alternative {i} missing 'index' key: {list(alt.keys())}"
        )
        assert "description" in alt, (
            f"Alternative {i} missing 'description' key: {list(alt.keys())}"
        )
        assert "chosen" in alt, (
            f"Alternative {i} missing 'chosen' key: {list(alt.keys())}"
        )


@then("exactly one alternative should have chosen set to true")
def step_exactly_one_chosen(context: Context) -> None:
    alts = context.pe_explain_dict["alternatives"]
    chosen_count = sum(1 for alt in alts if alt.get("chosen") is True)
    assert chosen_count == 1, (
        f"Expected exactly 1 chosen alternative, got {chosen_count}"
    )


@then('the json output should contain "{text}"')
def step_json_contains(context: Context, text: str) -> None:
    assert text in context.pe_json_output, f"Expected '{text}' in JSON output"


@then("the plan explain json output should be valid")
def step_json_valid(context: Context) -> None:
    parsed = json.loads(context.pe_json_output)
    assert isinstance(parsed, dict), "Expected a JSON object"


@then('the yaml output should contain "{text}"')
def step_yaml_contains(context: Context, text: str) -> None:
    assert text in context.pe_yaml_output, f"Expected '{text}' in YAML output"


@then("the explain lookup should indicate not found")
def step_explain_not_found(context: Context) -> None:
    # Verify id is a valid ULID
    assert context.pe_missing_id is not None
    assert len(context.pe_missing_id) == 26  # ULID length
    # Verify the generated id is NOT among any known decision ids
    assert context.pe_missing_id not in context.pe_known_ids, (
        "Non-existent id should differ from all known decision ids"
    )
    # _build_explain_dict requires a Decision object; passing None would be
    # a TypeError.  The "not found" path is handled at the CLI level
    # (explain_decision_cmd checks `decision is None` before calling
    # _build_explain_dict).  Verify the contract: the function signature
    # does not accept None.
    import inspect

    sig = inspect.signature(_build_explain_dict)
    first_param = next(iter(sig.parameters.values()))
    assert first_param.annotation is not None


# ---------------------------------------------------------------------------
# Then steps - tree
# ---------------------------------------------------------------------------


@then("the tree should have at least {count:d} root node")
def step_tree_root_count(context: Context, count: int) -> None:
    assert len(context.pe_tree) >= count, (
        f"Expected >= {count} roots, got {len(context.pe_tree)}"
    )


@then("the first root node should have children")
def step_root_has_children(context: Context) -> None:
    root = context.pe_tree[0]
    assert len(root["children"]) > 0, "Expected root to have children"


@then("the tree should include superseded decision nodes")
def step_tree_includes_superseded(context: Context) -> None:
    # With show_superseded=True, all decisions including superseded should be in tree
    all_ids: list[str] = []
    _collect_ids(context.pe_tree, all_ids)
    assert context.pe_superseded_id in all_ids, (
        f"Superseded ID {context.pe_superseded_id} not found in tree"
    )


@then("the root nodes should have no children")
def step_roots_no_children(context: Context) -> None:
    for root in context.pe_tree:
        assert len(root["children"]) == 0, (
            f"Root {root['decision_id']} should have no children at depth 1"
        )


@then("the json tree output should be valid json")
def step_tree_json_valid(context: Context) -> None:
    parsed = json.loads(context.pe_tree_json)
    assert isinstance(parsed, dict), "Expected a JSON envelope object"
    assert "data" in parsed, "Expected data key in JSON envelope"
    # format_output wraps data in a spec-required envelope dict;
    # the actual tree array is in the "data" key.
    assert isinstance(parsed["data"], list), (
        "Expected data to be a JSON array of tree nodes"
    )


@then("the json tree output should be a valid json envelope")
def step_tree_json_valid_envelope(context: Context) -> None:
    parsed = json.loads(context.pe_tree_json)
    assert isinstance(parsed, dict), (
        f"Expected a JSON object (envelope), got {type(parsed)}"
    )
    _envelope_keys = {"command", "status", "exit_code", "data", "timing", "messages"}
    assert _envelope_keys.issubset(parsed.keys()), (
        f"Expected envelope keys {_envelope_keys}, got {set(parsed.keys())}"
    )


@then('the json tree output should contain "{text}"')
def step_tree_json_contains(context: Context, text: str) -> None:
    assert text in context.pe_tree_json, f"Expected '{text}' in tree JSON"


@then('the yaml tree output should contain "{text}"')
def step_tree_yaml_contains(context: Context, text: str) -> None:
    assert text in context.pe_tree_yaml, f"Expected '{text}' in tree YAML"


@then("the tree should be empty")
def step_tree_empty(context: Context) -> None:
    assert len(context.pe_tree) == 0, (
        f"Expected empty tree, got {len(context.pe_tree)} roots"
    )


@then("the tree should not include superseded decision nodes")
def step_tree_excludes_superseded(context: Context) -> None:
    all_ids: list[str] = []
    _collect_ids(context.pe_tree, all_ids)
    assert context.pe_superseded_id not in all_ids, (
        f"Superseded ID {context.pe_superseded_id} should not be in default tree"
    )


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _collect_ids(nodes: list[dict[str, object]], out: list[str]) -> None:
    """Recursively collect all decision_ids from tree nodes."""
    from collections import deque

    queue: deque[dict[str, object]] = deque(nodes)
    while queue:
        node = queue.popleft()
        out.append(str(node["decision_id"]))
        children = node.get("children", [])
        assert isinstance(children, list)
        for child in children:
            assert isinstance(child, dict)
            queue.append(child)


# ---------------------------------------------------------------------------
# Given steps - per-type ordinals
# ---------------------------------------------------------------------------


@given("a set of test decisions with multiple types for ordinal testing")
def step_decisions_multiple_types(context: Context) -> None:
    """Create decisions with multiple invariants, strategies, and parallel spawns."""
    root_id = str(ULID())
    inv1_id = str(ULID())
    inv2_id = str(ULID())
    strat1_id = str(ULID())
    par1_id = str(ULID())

    context.pe_decisions = [
        Decision(
            decision_id=root_id,
            plan_id=_PLAN_ID,
            sequence_number=0,
            decision_type=DecisionType.PROMPT_DEFINITION,
            question="What to build?",
            chosen_option="REST API",
        ),
        Decision(
            decision_id=inv1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=1,
            decision_type=DecisionType.INVARIANT_ENFORCED,
            question="Enforce authentication?",
            chosen_option="OAuth2",
        ),
        Decision(
            decision_id=inv2_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=2,
            decision_type=DecisionType.INVARIANT_ENFORCED,
            question="Enforce rate limiting?",
            chosen_option="100 requests per minute",
        ),
        Decision(
            decision_id=strat1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=3,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question="Which strategy?",
            chosen_option="Microservices",
        ),
        Decision(
            decision_id=par1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=4,
            decision_type=DecisionType.SUBPLAN_PARALLEL_SPAWN,
            question="Run in parallel?",
            chosen_option="Yes",
        ),
    ]


@given("a set of test decisions with mixed types")
def step_decisions_mixed_types(context: Context) -> None:
    """Create decisions with invariants, strategies, and spawns."""
    root_id = str(ULID())
    inv1_id = str(ULID())
    inv2_id = str(ULID())
    strat1_id = str(ULID())
    strat2_id = str(ULID())
    spawn1_id = str(ULID())
    spawn2_id = str(ULID())

    context.pe_decisions = [
        Decision(
            decision_id=root_id,
            plan_id=_PLAN_ID,
            sequence_number=0,
            decision_type=DecisionType.PROMPT_DEFINITION,
            question="What to build?",
            chosen_option="REST API",
        ),
        Decision(
            decision_id=inv1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=1,
            decision_type=DecisionType.INVARIANT_ENFORCED,
            question="Enforce auth?",
            chosen_option="OAuth2",
        ),
        Decision(
            decision_id=inv2_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=2,
            decision_type=DecisionType.INVARIANT_ENFORCED,
            question="Enforce encryption?",
            chosen_option="TLS",
        ),
        Decision(
            decision_id=strat1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=3,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question="Which architecture?",
            chosen_option="Microservices",
        ),
        Decision(
            decision_id=strat2_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=4,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question="Which deployment?",
            chosen_option="Kubernetes",
        ),
        Decision(
            decision_id=spawn1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=5,
            decision_type=DecisionType.SUBPLAN_SPAWN,
            question="Spawn auth service?",
            chosen_option="Yes",
        ),
        Decision(
            decision_id=spawn2_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=6,
            decision_type=DecisionType.SUBPLAN_SPAWN,
            question="Spawn payment service?",
            chosen_option="Yes",
        ),
    ]


@given("a set of test decisions forming a tree with multiple types")
def step_tree_multiple_types(context: Context) -> None:
    """Create a complete tree with different decision types."""
    root_id = str(ULID())
    inv1_id = str(ULID())
    inv2_id = str(ULID())
    strat1_id = str(ULID())
    par1_id = str(ULID())

    context.pe_decisions = [
        Decision(
            decision_id=root_id,
            plan_id=_PLAN_ID,
            sequence_number=0,
            decision_type=DecisionType.PROMPT_DEFINITION,
            question="What to build?",
            chosen_option="REST API",
        ),
        Decision(
            decision_id=inv1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=1,
            decision_type=DecisionType.INVARIANT_ENFORCED,
            question="Enforce authentication?",
            chosen_option="OAuth2",
        ),
        Decision(
            decision_id=inv2_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=2,
            decision_type=DecisionType.INVARIANT_ENFORCED,
            question="Enforce rate limiting?",
            chosen_option="100 req/min",
        ),
        Decision(
            decision_id=strat1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=3,
            decision_type=DecisionType.STRATEGY_CHOICE,
            question="Which framework?",
            chosen_option="FastAPI",
        ),
        Decision(
            decision_id=par1_id,
            plan_id=_PLAN_ID,
            parent_decision_id=root_id,
            sequence_number=4,
            decision_type=DecisionType.SUBPLAN_PARALLEL_SPAWN,
            question="Run in parallel?",
            chosen_option="Yes",
        ),
    ]


# ---------------------------------------------------------------------------
# When steps - per-type ordinals
# ---------------------------------------------------------------------------


@when("I generate decision labels with per-type ordinals")
def step_generate_labels(context: Context) -> None:
    """Generate labels for decisions using per-type ordinals."""
    # Build the per-type ordinal counter (same logic as in the CLI)
    type_counts: dict[str, int] = {}
    decision_labels: dict[str, str] = {}

    for d in context.pe_decisions:
        type_counts[d.decision_type] = type_counts.get(d.decision_type, 0) + 1
        per_type_ordinal = type_counts[d.decision_type]
        decision_labels[d.decision_id] = _get_decision_label(
            d.decision_type, per_type_ordinal
        )

    context.pe_decision_labels = decision_labels


@when("I build the plain format tree output with decision labels")
def step_build_tree_with_labels(context: Context) -> None:
    """Build tree output and capture the decision labels section."""
    # Build the tree
    filtered = context.pe_decisions

    # Build per-type ordinals (same logic as in the CLI command)
    type_counts: dict[str, int] = {}
    decision_ordinals: dict[str, int] = {}
    for d in filtered:
        type_counts[d.decision_type] = type_counts.get(d.decision_type, 0) + 1
        decision_ordinals[d.decision_id] = type_counts[d.decision_type]

    # Capture the decision labels output
    labels_section = []
    labels_section.append("Decision IDs (for correction)")
    for d in filtered:
        type_label = _get_decision_label(
            d.decision_type, decision_ordinals[d.decision_id]
        )
        labels_section.append(f"  {type_label}: {d.decision_id}")

    context.pe_tree_output = "\n".join(labels_section)


# ---------------------------------------------------------------------------
# Then steps - per-type ordinals
# ---------------------------------------------------------------------------


@then('the first invariant should be labeled "{label}"')
def step_first_invariant_label(context: Context, label: str) -> None:
    """Verify the first invariant has the correct label."""
    inv_ids = [
        d.decision_id
        for d in context.pe_decisions
        if d.decision_type == DecisionType.INVARIANT_ENFORCED
    ]
    assert len(inv_ids) > 0, "No invariant decisions found"
    first_inv = inv_ids[0]
    assert context.pe_decision_labels[first_inv] == label, (
        f"Expected first invariant to be '{label}', "
        f"got '{context.pe_decision_labels[first_inv]}'"
    )


@then('the second invariant should be labeled "{label}"')
def step_second_invariant_label(context: Context, label: str) -> None:
    """Verify the second invariant has the correct label."""
    inv_ids = [
        d.decision_id
        for d in context.pe_decisions
        if d.decision_type == DecisionType.INVARIANT_ENFORCED
    ]
    assert len(inv_ids) > 1, "Fewer than 2 invariant decisions found"
    second_inv = inv_ids[1]
    assert context.pe_decision_labels[second_inv] == label, (
        f"Expected second invariant to be '{label}', "
        f"got '{context.pe_decision_labels[second_inv]}'"
    )


@then('the first strategy should be labeled "{label}"')
def step_first_strategy_label(context: Context, label: str) -> None:
    """Verify the first strategy has the correct label (no ordinal)."""
    strat_ids = [
        d.decision_id
        for d in context.pe_decisions
        if d.decision_type == DecisionType.STRATEGY_CHOICE
    ]
    assert len(strat_ids) > 0, "No strategy decisions found"
    first_strat = strat_ids[0]
    actual_label = context.pe_decision_labels[first_strat]
    assert actual_label == label, (
        f"Expected first strategy to be '{label}', got '{actual_label}'"
    )


@then('the first parallel should be labeled "{label}"')
def step_first_parallel_label(context: Context, label: str) -> None:
    """Verify the first parallel spawn has the correct label."""
    par_ids = [
        d.decision_id
        for d in context.pe_decisions
        if d.decision_type == DecisionType.SUBPLAN_PARALLEL_SPAWN
    ]
    assert len(par_ids) > 0, "No parallel spawn decisions found"
    first_par = par_ids[0]
    assert context.pe_decision_labels[first_par] == label, (
        f"Expected first parallel to be '{label}', "
        f"got '{context.pe_decision_labels[first_par]}'"
    )


@then("each decision type should start counting from 1")
def step_each_type_starts_from_one(context: Context) -> None:
    """Verify each decision type starts counting from 1."""
    # Check that for each decision type, the ordinal numbers are sequential
    type_ordinals: dict[str, list[int]] = {}

    for d in context.pe_decisions:
        label = context.pe_decision_labels[d.decision_id]
        # Extract the ordinal number from the label (e.g., "Invariant 2" -> 2)
        parts = label.split()
        if len(parts) > 1 and parts[-1].isdigit():
            ordinal = int(parts[-1])
            type_ordinals.setdefault(d.decision_type, []).append(ordinal)

    # For each type with ordinals, verify they start from 1
    for dtype, ordinals in type_ordinals.items():
        sorted_ordinals = sorted(set(ordinals))
        expected = list(range(1, len(sorted_ordinals) + 1))
        assert sorted_ordinals == expected, (
            f"Decision type {dtype} ordinals {sorted_ordinals} != expected {expected}"
        )


@then("invariants should have sequential numbers")
def step_invariants_sequential(context: Context) -> None:
    """Verify invariant ordinals are sequential."""
    inv_labels = [
        context.pe_decision_labels[d.decision_id]
        for d in context.pe_decisions
        if d.decision_type == DecisionType.INVARIANT_ENFORCED
    ]
    # Extract ordinals from labels
    ordinals = []
    for label in inv_labels:
        parts = label.split()
        if len(parts) > 1 and parts[-1].isdigit():
            ordinals.append(int(parts[-1]))
    # Should be [1, 2, ...] in order
    expected = list(range(1, len(ordinals) + 1))
    assert ordinals == expected, f"Invariant ordinals {ordinals} != {expected}"


@then("parallel spawns should have sequential numbers")
def step_parallel_spawns_sequential(context: Context) -> None:
    """Verify parallel spawn ordinals are sequential."""
    par_labels = [
        context.pe_decision_labels[d.decision_id]
        for d in context.pe_decisions
        if d.decision_type == DecisionType.SUBPLAN_PARALLEL_SPAWN
    ]
    # Extract ordinals from labels
    ordinals = []
    for label in par_labels:
        parts = label.split()
        if len(parts) > 1 and parts[-1].isdigit():
            ordinals.append(int(parts[-1]))
    # Should be [1, 2, ...] in order
    expected = list(range(1, len(ordinals) + 1))
    assert ordinals == expected, f"Parallel ordinals {ordinals} != {expected}"


@then("spawn decisions should have sequential numbers")
def step_spawns_sequential(context: Context) -> None:
    """Verify spawn ordinals are sequential."""
    spawn_labels = [
        context.pe_decision_labels[d.decision_id]
        for d in context.pe_decisions
        if d.decision_type == DecisionType.SUBPLAN_SPAWN
    ]
    # Extract ordinals from labels
    ordinals = []
    for label in spawn_labels:
        parts = label.split()
        if len(parts) > 1 and parts[-1].isdigit():
            ordinals.append(int(parts[-1]))
    # Should be [1, 2, ...] in order
    expected = list(range(1, len(ordinals) + 1))
    assert ordinals == expected, f"Spawn ordinals {ordinals} != {expected}"


@then('the decision tree output should contain "{text}"')
def step_decision_tree_output_contains(context: Context, text: str) -> None:
    """Verify text is in decision tree output."""
    assert text in context.pe_tree_output, (
        f"Expected '{text}' in output:\n{context.pe_tree_output}"
    )


@then('the decision tree output should not contain "{text}"')
def step_decision_tree_output_not_contains(context: Context, text: str) -> None:
    """Verify text does not appear in decision tree output."""
    assert text not in context.pe_tree_output, (
        f"Unexpected '{text}' in output:\n{context.pe_tree_output}"
    )