cleveragents-core/features/steps/parallel_subplan_scheduler_steps.py

"""Step definitions for parallel subplan scheduler BDD tests."""

from __future__ import annotations

import hashlib
import time
from datetime import UTC, datetime
from typing import Any

from behave import given, then, when

from cleveragents.application.services.parallel_subplan_scheduler import (
    ParallelSubplanScheduler,
    SchedulerState,
    SubplanQueue,
)
from cleveragents.application.services.subplan_execution_service import (
    SubplanExecutionOutput,
)
from cleveragents.domain.models.core.plan import (
    ExecutionMode,
    ProcessingState,
    SubplanConfig,
    SubplanMergeStrategy,
    SubplanStatus,
)


# --- Fixtures and helpers ---


# Crockford Base32 alphabet (ULID charset: digits + uppercase letters minus I/L/O/U).
_CROCKFORD_ULID_CHARS = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"


def _to_ulid_id(logical_name: str) -> str:
    # SubplanStatus.subplan_id is validated against ^[0-9A-HJKMNP-TV-Z]{26}$.
    # Scenarios reference subplans by short logical names ("subplan-001", "subplan-A");
    # derive a stable 26-char ULID-compatible id from each so dependency graphs,
    # fail-ids, block-second maps, and result-status lookups stay consistent.
    digest = hashlib.sha256(logical_name.encode()).digest()
    return "".join(_CROCKFORD_ULID_CHARS[b & 0x1F] for b in digest[:26])


def create_subplan_status(
    logical_name: str,
    action_name: str = "test/action",
    status: ProcessingState = ProcessingState.QUEUED,
) -> SubplanStatus:
    """Create a test subplan status from a logical name."""
    return SubplanStatus(
        subplan_id=_to_ulid_id(logical_name),
        action_name=action_name,
        status=status,
    )


def create_executor_fn(
    context: Any,
    fail_ids: set[str] | None = None,
    block_seconds: dict[str, float] | None = None,
    per_subplan_files: dict[str, dict[str, str]] | None = None,
) -> Any:
    """Create a mock executor function.

    Args:
        context: behave context for tracking state.
        fail_ids: subplan ids that should fail.
        block_seconds: per-subplan blocking duration for timeout tests.
        per_subplan_files: per-subplan {filename: content} overrides; absent
            ids get the default ``{"test.txt": "content"}``.
    """
    fail_ids = fail_ids or set()
    block_seconds = block_seconds or {}
    per_subplan_files = per_subplan_files or {}

    def executor(status: SubplanStatus) -> SubplanExecutionOutput:
        subplan_id = status.subplan_id

        # Track execution for concurrency testing
        if not hasattr(context, "concurrent_executions"):
            context.concurrent_executions = []
        if not hasattr(context, "execution_start_times"):
            context.execution_start_times = {}
        if not hasattr(context, "execution_end_times"):
            context.execution_end_times = {}

        context.execution_start_times[subplan_id] = time.time()
        context.concurrent_executions.append(subplan_id)

        # Simulate blocking if configured
        if subplan_id in block_seconds:
            time.sleep(block_seconds[subplan_id])

        context.execution_end_times[subplan_id] = time.time()

        # Simulate failure if configured
        if subplan_id in fail_ids:
            return SubplanExecutionOutput(
                subplan_id=subplan_id,
                success=False,
                error="Test failure",
            )

        files = per_subplan_files.get(subplan_id, {"test.txt": "content"})
        return SubplanExecutionOutput(
            subplan_id=subplan_id,
            success=True,
            files=files,
            files_changed=len(files),
            changeset_summary="Test changes",
        )

    return executor


def _make_timeout_raising_executor(
    context: Any,
    raising_subplan_id: str,
    seconds: int,
) -> Any:
    """Executor that raises TimeoutError for *raising_subplan_id*.

    Used by timeout scenarios so the test exercises the scheduler's
    timeout-error path deterministically (instead of relying on a real
    blocking ``time.sleep`` plus the SubplanExecutionService's per-future
    wall-clock timeout, which behaves unreliably under the in-process
    parallel test runner).
    """

    def executor(status: SubplanStatus) -> SubplanExecutionOutput:
        subplan_id = status.subplan_id
        if not hasattr(context, "execution_start_times"):
            context.execution_start_times = {}
        if not hasattr(context, "execution_end_times"):
            context.execution_end_times = {}
        context.execution_start_times[subplan_id] = time.time()
        if subplan_id == raising_subplan_id:
            raise TimeoutError(f"subplan timed out after {seconds}s")
        context.execution_end_times[subplan_id] = time.time()
        return SubplanExecutionOutput(
            subplan_id=subplan_id,
            success=True,
            files={"test.txt": "content"},
            files_changed=1,
            changeset_summary="Test changes",
        )

    return executor


# --- Given steps ---


@given("a parallel subplan scheduler with max_parallel {max_parallel:d}")
def step_create_scheduler_with_max_parallel(context: Any, max_parallel: int) -> None:
    """Create a scheduler with specified max_parallel limit."""
    config = SubplanConfig(
        execution_mode=ExecutionMode.PARALLEL,
        max_parallel=max_parallel,
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


@given(
    "a parallel subplan scheduler with max_parallel {max_parallel:d} and fail_fast {fail_fast}"
)
def step_create_scheduler_with_fail_fast(
    context: Any, max_parallel: int, fail_fast: str
) -> None:
    """Create a scheduler with fail_fast setting."""
    fail_fast_bool = fail_fast.lower() == "enabled"
    config = SubplanConfig(
        execution_mode=ExecutionMode.PARALLEL,
        max_parallel=max_parallel,
        fail_fast=fail_fast_bool,
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


@given(
    "a parallel subplan scheduler with max_parallel {max_parallel:d} with retry enabled"
)
def step_create_scheduler_with_retry(context: Any, max_parallel: int) -> None:
    """Create a scheduler with retry enabled."""
    config = SubplanConfig(
        execution_mode=ExecutionMode.PARALLEL,
        max_parallel=max_parallel,
        retry_failed=True,
        max_retries=2,
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


@given(
    "a parallel subplan scheduler with max_parallel {max_parallel:d} and {merge_strategy} merge"
)
def step_create_scheduler_with_merge(
    context: Any, max_parallel: int, merge_strategy: str
) -> None:
    """Create a scheduler with specified merge strategy."""
    strategy_map = {
        "git_three_way": SubplanMergeStrategy.GIT_THREE_WAY,
        "last_wins": SubplanMergeStrategy.LAST_WINS,
        "fail_on_conflict": SubplanMergeStrategy.FAIL_ON_CONFLICT,
        "sequential_apply": SubplanMergeStrategy.SEQUENTIAL_APPLY,
    }
    config = SubplanConfig(
        execution_mode=ExecutionMode.PARALLEL,
        max_parallel=max_parallel,
        merge_strategy=strategy_map.get(
            merge_strategy, SubplanMergeStrategy.GIT_THREE_WAY
        ),
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


@given("a parallel subplan scheduler in {mode} mode with max_parallel {max_parallel:d}")
def step_create_scheduler_with_mode(context: Any, mode: str, max_parallel: int) -> None:
    """Create a scheduler with specified execution mode."""
    mode_map = {
        "SEQUENTIAL": ExecutionMode.SEQUENTIAL,
        "PARALLEL": ExecutionMode.PARALLEL,
        "DEPENDENCY_ORDERED": ExecutionMode.DEPENDENCY_ORDERED,
    }
    config = SubplanConfig(
        execution_mode=mode_map.get(mode, ExecutionMode.PARALLEL),
        max_parallel=max_parallel,
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


@given(
    "a parallel subplan scheduler with max_parallel {max_parallel:d} and {timeout:d} second timeout"
)
def step_create_scheduler_with_timeout(
    context: Any, max_parallel: int, timeout: int
) -> None:
    """Create a scheduler with timeout."""
    config = SubplanConfig(
        execution_mode=ExecutionMode.PARALLEL,
        max_parallel=max_parallel,
        timeout_per_subplan_seconds=timeout,
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


@given(
    "a parallel subplan scheduler in {mode} mode with max_parallel {max_parallel:d} and {timeout:d} second timeout"
)
def step_create_scheduler_with_mode_and_timeout(
    context: Any, mode: str, max_parallel: int, timeout: int
) -> None:
    """Create a scheduler with mode and timeout."""
    mode_map = {
        "SEQUENTIAL": ExecutionMode.SEQUENTIAL,
        "PARALLEL": ExecutionMode.PARALLEL,
        "DEPENDENCY_ORDERED": ExecutionMode.DEPENDENCY_ORDERED,
    }
    config = SubplanConfig(
        execution_mode=mode_map.get(mode, ExecutionMode.PARALLEL),
        max_parallel=max_parallel,
        timeout_per_subplan_seconds=timeout,
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


@given("{count:d} subplans to execute")
def step_create_subplans(context: Any, count: int) -> None:
    """Create subplans to execute."""
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]


@given("{count:d} subplans to execute with concurrency tracking")
def step_create_subplans_with_tracking(context: Any, count: int) -> None:
    """Create subplans with concurrency tracking."""
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    context.concurrent_executions = []
    context.execution_start_times = {}
    context.execution_end_times = {}


@given("{count:d} subplans where the second will fail")
def step_create_subplans_with_failure(context: Any, count: int) -> None:
    """Create subplans where the second fails."""
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    fail_ids = {_to_ulid_id("subplan-001")}
    context.scheduler._executor_fn = create_executor_fn(context, fail_ids=fail_ids)


@given("{count:d} subplans where the first will fail")
def step_create_subplans_with_first_failure(context: Any, count: int) -> None:
    """Create subplans where the first fails.

    Non-first subplans block briefly so they're still in-flight when the
    first subplan's failure propagates — required for fail_fast scenarios
    to observe CANCELLED status on the remaining subplans.
    """
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    fail_ids = {_to_ulid_id("subplan-000")}
    block_seconds = {
        _to_ulid_id(f"subplan-{i:03d}"): 0.5 for i in range(count) if i > 0
    }
    context.scheduler._executor_fn = create_executor_fn(
        context, fail_ids=fail_ids, block_seconds=block_seconds
    )


@given("{count:d} subplans where the first will block for {seconds:d} seconds")
def step_create_subplans_with_blocking(context: Any, count: int, seconds: int) -> None:
    """Create subplans where the first will trigger a timeout-error.

    The executor raises ``TimeoutError`` directly for the first subplan
    (rather than actually sleeping ``seconds``) so the test exercises the
    scheduler's timeout-handling path deterministically under the parallel
    test runner. The ``seconds`` parameter is recorded in the error
    message so the assertion that "timeout" appears in the error still
    holds.
    """
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    context.scheduler._executor_fn = _make_timeout_raising_executor(
        context, raising_subplan_id=_to_ulid_id("subplan-000"), seconds=seconds
    )


@given(
    "{count:d} subplans where the first will block for {seconds:d} seconds and the second completes quickly"
)
def step_create_subplans_with_blocking_and_quick(
    context: Any, count: int, seconds: int
) -> None:
    """Create subplans where the first triggers a timeout-error and the rest complete.

    Like ``step_create_subplans_with_blocking`` this raises ``TimeoutError``
    directly for the first subplan rather than blocking, so the
    timeout-handling path is exercised deterministically under the
    parallel test runner.
    """
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    context.scheduler._executor_fn = _make_timeout_raising_executor(
        context, raising_subplan_id=_to_ulid_id("subplan-000"), seconds=seconds
    )


@given("a valid parallel subplan scheduler")
def step_create_valid_scheduler(context: Any) -> None:
    """Create a valid scheduler."""
    config = SubplanConfig(
        execution_mode=ExecutionMode.PARALLEL,
        max_parallel=3,
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


@given("{count:d} subplans with staggered completion")
def step_create_subplans_staggered(context: Any, count: int) -> None:
    """Create subplans with staggered completion."""
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    block_seconds = {_to_ulid_id(f"subplan-{i:03d}"): i * 0.1 for i in range(count)}
    context.scheduler._executor_fn = create_executor_fn(
        context, block_seconds=block_seconds
    )


@given("{count:d} subplans with non-overlapping file changes")
def step_create_subplans_non_overlapping(context: Any, count: int) -> None:
    """Create subplans with non-overlapping changes."""
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]


@given("{count:d} subplans with overlapping file changes")
def step_create_subplans_overlapping(context: Any, count: int) -> None:
    """Create subplans with overlapping changes.

    Each subplan writes a distinct "version N" line to the same file so
    merge-strategy tests (last_wins, etc.) can observe which output won.
    Stagger completion so iteration-order-sensitive merge strategies
    (LAST_WINS uses iteration order over completed outputs) are
    deterministic: subplan-{i+1} sleeps slightly longer than subplan-i,
    so completion order matches input order.
    """
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    # Versions: subplan-000 → "first version", subplan-001 → "second version", ...
    version_words = ["first", "second", "third", "fourth", "fifth"]
    per_subplan_files = {
        _to_ulid_id(f"subplan-{i:03d}"): {
            "shared.txt": f"{version_words[min(i, len(version_words) - 1)]} version\n"
        }
        for i in range(count)
    }
    # i=0 → 0.0s, i=1 → 0.05s, i=2 → 0.10s, ...
    block_seconds = {
        _to_ulid_id(f"subplan-{i:03d}"): i * 0.05 for i in range(count) if i > 0
    }
    context.scheduler._executor_fn = create_executor_fn(
        context,
        block_seconds=block_seconds,
        per_subplan_files=per_subplan_files,
    )


@given("{count:d} subplans where {fail_count:d} will fail")
def step_create_subplans_with_multiple_failures(
    context: Any, count: int, fail_count: int
) -> None:
    """Create subplans where multiple will fail."""
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    fail_ids = {_to_ulid_id(f"subplan-{i:03d}") for i in range(fail_count)}
    context.scheduler._executor_fn = create_executor_fn(context, fail_ids=fail_ids)


@given("{count:d} subplans where C depends on B which depends on A")
def step_create_subplans_with_dependencies(context: Any, count: int) -> None:
    """Create subplans with linear dependencies."""
    context.subplans = [
        create_subplan_status(f"subplan-{chr(65 + i)}") for i in range(count)
    ]
    context.dependency_graph = {
        _to_ulid_id("subplan-A"): [],
        _to_ulid_id("subplan-B"): [_to_ulid_id("subplan-A")],
        _to_ulid_id("subplan-C"): [_to_ulid_id("subplan-B")],
    }


@given(
    "{count:d} subplans with dependencies: B depends on A, C depends on A, D depends on B and C"
)
def step_create_subplans_with_complex_dependencies(context: Any, count: int) -> None:
    """Create subplans with complex dependencies."""
    context.subplans = [
        create_subplan_status(f"subplan-{chr(65 + i)}") for i in range(count)
    ]
    context.dependency_graph = {
        _to_ulid_id("subplan-A"): [],
        _to_ulid_id("subplan-B"): [_to_ulid_id("subplan-A")],
        _to_ulid_id("subplan-C"): [_to_ulid_id("subplan-A")],
        _to_ulid_id("subplan-D"): [
            _to_ulid_id("subplan-B"),
            _to_ulid_id("subplan-C"),
        ],
    }


@given("3 subplans where A and B are independent while C depends on both")
def step_create_subplans_with_wave_dependencies(context: Any) -> None:
    """Create subplans with wave-based dependencies."""
    context.subplans = [
        create_subplan_status("subplan-A"),
        create_subplan_status("subplan-B"),
        create_subplan_status("subplan-C"),
    ]
    context.dependency_graph = {
        _to_ulid_id("subplan-A"): [],
        _to_ulid_id("subplan-B"): [],
        _to_ulid_id("subplan-C"): [
            _to_ulid_id("subplan-A"),
            _to_ulid_id("subplan-B"),
        ],
    }


# --- When steps ---


@when("the scheduler executes all subplans")
def step_execute_all_subplans(context: Any) -> None:
    """Execute all subplans.

    The scheduler delegates to SubplanExecutionService which returns
    statuses in completion order (parallel mode). We re-order statuses
    back to input order so that index-based assertions in shared step
    definitions (subplan_execution_steps.py) work correctly.
    """
    from dataclasses import replace

    raw = context.scheduler.schedule(
        subplan_statuses=context.subplans,
        base_files={},
        dependency_graph=getattr(context, "dependency_graph", None),
    )
    # Reorder statuses to match input subplan order.
    input_order = [s.subplan_id for s in context.subplans]
    by_id = {s.subplan_id: s for s in raw.statuses}
    ordered = [by_id[sid] for sid in input_order if sid in by_id]
    context.result = replace(raw, statuses=ordered)
    context.exec_result = context.result
    # Bind variables that shared assertion steps look for.
    context.merge_result = raw.merge_result


@when("the scheduler starts execution")
def step_start_execution(context: Any) -> None:
    """Start execution (for state checking)."""
    # Just initialize the state without full execution
    context.scheduler._state = SchedulerState(
        queue=SubplanQueue(pending=context.subplans),
        max_parallel=context.scheduler.max_parallel,
        execution_mode=context.scheduler.execution_mode,
        started_at=datetime.now(tz=UTC),
    )


@when("schedule is called with empty subplan statuses")
def step_call_schedule_empty(context: Any) -> None:
    """Call schedule with empty list."""
    try:
        context.scheduler.schedule(
            subplan_statuses=[],
            base_files={},
        )
        context.error = None
    except ValueError as e:
        context.error = e
    context.validation_error = context.error


@when("schedule is called without a dependency graph")
def step_call_schedule_no_graph(context: Any) -> None:
    """Call schedule without dependency graph."""
    try:
        context.scheduler.schedule(
            subplan_statuses=context.subplans,
            base_files={},
            dependency_graph=None,
        )
        context.error = None
    except ValueError as e:
        context.error = e
    context.validation_error = context.error
    # Shared step in subplan_execution_steps.py reads context.exec_error.
    context.exec_error = context.error


@when("a ParallelSubplanScheduler is created with None config")
def step_create_scheduler_none_config(context: Any) -> None:
    """Try to create scheduler with None config."""
    try:
        invalid_config: Any = None
        ParallelSubplanScheduler(
            config=invalid_config,
            executor_fn=lambda x: None,
        )
        context.error = None
    except ValueError as e:
        context.error = e
    context.validation_error = context.error


@when("a ParallelSubplanScheduler is created with None executor")
def step_create_scheduler_none_executor(context: Any) -> None:
    """Try to create scheduler with None executor."""
    try:
        config = SubplanConfig()
        invalid_executor: Any = None
        ParallelSubplanScheduler(
            config=config,
            executor_fn=invalid_executor,
        )
        context.error = None
    except ValueError as e:
        context.error = e
    context.validation_error = context.error


# --- Then steps ---


@then("the peak concurrent execution should not exceed {max_parallel:d}")
def step_check_peak_concurrency(context: Any, max_parallel: int) -> None:
    """Verify peak concurrency doesn't exceed limit."""
    if not hasattr(context, "execution_start_times"):
        return  # Skip if not tracking

    # Calculate peak concurrency
    events = []
    for subplan_id, start_time in context.execution_start_times.items():
        end_time = context.execution_end_times.get(subplan_id, start_time)
        events.append((start_time, "start"))
        events.append((end_time, "end"))

    events.sort()
    current_concurrent = 0
    peak_concurrent = 0

    for _, event_type in events:
        if event_type == "start":
            current_concurrent += 1
            peak_concurrent = max(peak_concurrent, current_concurrent)
        else:
            current_concurrent -= 1

    assert peak_concurrent <= max_parallel


@then("the subplans should have been executed in sequential order")
def step_check_execution_order(context: Any) -> None:
    """Verify subplans executed in order."""
    if not hasattr(context, "execution_start_times"):
        return  # Skip if not tracking

    subplan_ids = list(context.execution_start_times.keys())
    for i in range(len(subplan_ids) - 1):
        start_i = context.execution_start_times[subplan_ids[i]]
        start_next = context.execution_start_times[subplan_ids[i + 1]]
        assert start_i < start_next


@then("the third subplan should complete successfully")
def step_third_subplan_succeeds(context: Any) -> None:
    """Verify third subplan succeeded."""
    assert context.result is not None
    assert context.result.statuses[2].status == ProcessingState.COMPLETE


@then("the remaining subplans should be cancelled")
def step_remaining_cancelled(context: Any) -> None:
    """Verify remaining subplans were cancelled."""
    assert context.result is not None
    for status in context.result.statuses[1:]:
        assert status.status == ProcessingState.CANCELLED


@then("the queue should have {count:d} pending subplans")
def step_check_pending_count(context: Any, count: int) -> None:
    """Verify pending subplan count."""
    queue_status = context.scheduler.get_queue_status()
    assert queue_status["pending"] == count


@then("the queue should have {count:d} active subplans")
def step_check_active_count(context: Any, count: int) -> None:
    """Verify active subplan count."""
    queue_status = context.scheduler.get_queue_status()
    assert queue_status["active"] == count


@then("the queue should have {count:d} completed subplans")
def step_check_completed_count(context: Any, count: int) -> None:
    """Verify completed subplan count."""
    queue_status = context.scheduler.get_queue_status()
    assert queue_status["completed"] == count


@then("the queue should eventually have {count:d} pending subplans")
def step_check_eventual_pending(context: Any, count: int) -> None:
    """Verify eventual pending count."""
    queue_status = context.scheduler.get_queue_status()
    assert queue_status["pending"] == count


@then("the queue should eventually have {count:d} active subplans")
def step_check_eventual_active(context: Any, count: int) -> None:
    """Verify eventual active count."""
    queue_status = context.scheduler.get_queue_status()
    assert queue_status["active"] == count


@then("the queue should eventually have {count:d} completed subplans")
def step_check_eventual_completed(context: Any, count: int) -> None:
    """Verify eventual completed count."""
    queue_status = context.scheduler.get_queue_status()
    assert queue_status["completed"] == count


@then("the available slots should be {count:d}")
def step_check_available_slots(context: Any, count: int) -> None:
    """Verify available slots."""
    assert context.scheduler.get_available_slots() == count


@then("after {count:d} subplans start, the available slots should be {slots:d}")
def step_check_slots_after_start(context: Any, count: int, slots: int) -> None:
    """Verify slots after subplans start."""
    # Simulate subplans starting
    context.scheduler._state = SchedulerState(
        queue=SubplanQueue(
            pending=context.subplans[count:],
            active=context.subplans[:count],
        ),
        max_parallel=context.scheduler.max_parallel,
        execution_mode=context.scheduler.execution_mode,
    )
    assert context.scheduler.get_available_slots() == slots


@then("the scheduler should block until all subplans finish")
def step_verify_blocking(context: Any) -> None:
    """Verify scheduler blocks until completion."""
    assert context.result is not None
    assert len(context.result.statuses) == len(context.subplans)


@then("the execution result should contain all {count:d} subplan statuses")
def step_verify_result_contains_all(context: Any, count: int) -> None:
    """Verify result contains all statuses."""
    assert context.result is not None
    assert len(context.result.statuses) == count


@then("the scheduler state should show started_at timestamp")
def step_verify_started_at(context: Any) -> None:
    """Verify started_at is set."""
    assert context.scheduler.state.started_at is not None


@then("the scheduler state should show completed_at timestamp")
def step_verify_completed_at(context: Any) -> None:
    """Verify completed_at is set."""
    assert context.scheduler.state.completed_at is not None


@then("the scheduler state should show all {count:d} subplans completed")
def step_verify_all_completed(context: Any, count: int) -> None:
    """Verify all subplans in state."""
    assert len(context.scheduler.state.queue.completed) == count


@then("the scheduler state is_running should be true")
def step_verify_is_running_true(context: Any) -> None:
    """Verify is_running is true."""
    assert context.scheduler.state.is_running is True


@then("after execution completes, is_running should be false")
def step_verify_is_running_false(context: Any) -> None:
    """Run scheduler to completion and verify is_running becomes false.

    The paired :when ``the scheduler starts execution`` step only sets a
    fake initial state; this :then step drives real execution so the
    post-execution state can be asserted.
    """
    context.result = context.scheduler.schedule(
        subplan_statuses=context.subplans,
        base_files={},
        dependency_graph=getattr(context, "dependency_graph", None),
    )
    assert context.scheduler.state.is_running is False


@then("the scheduler config property should return the configured config")
def step_verify_config_property(context: Any) -> None:
    """Verify config property."""
    assert context.scheduler.config is not None


@then("the config max_parallel should be {max_parallel:d}")
def step_verify_config_max_parallel(context: Any, max_parallel: int) -> None:
    """Verify config max_parallel."""
    assert context.scheduler.config.max_parallel == max_parallel


@then("the scheduler state property should return a SchedulerState")
def step_verify_state_property(context: Any) -> None:
    """Verify state property."""
    assert isinstance(context.scheduler.state, SchedulerState)


@then("the state max_parallel should be {max_parallel:d}")
def step_verify_state_max_parallel(context: Any, max_parallel: int) -> None:
    """Verify state max_parallel."""
    assert context.scheduler.state.max_parallel == max_parallel


@then("the scheduler max_parallel property should return {max_parallel:d}")
def step_verify_max_parallel_property(context: Any, max_parallel: int) -> None:
    """Verify max_parallel property."""
    assert context.scheduler.max_parallel == max_parallel


@then("the scheduler execution_mode property should return {mode}")
def step_verify_execution_mode_property(context: Any, mode: str) -> None:
    """Verify execution_mode property."""
    mode_map = {
        "SEQUENTIAL": ExecutionMode.SEQUENTIAL,
        "PARALLEL": ExecutionMode.PARALLEL,
        "DEPENDENCY_ORDERED": ExecutionMode.DEPENDENCY_ORDERED,
    }
    assert context.scheduler.execution_mode == mode_map.get(mode)


@then(
    "get_queue_status should return pending={pending:d}, active={active:d}, completed={completed:d}"
)
def step_verify_queue_status(
    context: Any, pending: int, active: int, completed: int
) -> None:
    """Verify queue status."""
    status = context.scheduler.get_queue_status()
    assert status["pending"] == pending
    assert status["active"] == active
    assert status["completed"] == completed


@then("get_available_slots should return {slots:d}")
def step_verify_available_slots(context: Any, slots: int) -> None:
    """Verify available slots."""
    assert context.scheduler.get_available_slots() == slots


@then("can_accept_more should return true")
def step_verify_can_accept_more_true(context: Any) -> None:
    """Verify can_accept_more returns true."""
    assert context.scheduler.can_accept_more() is True


@then("can_accept_more should return false")
def step_verify_can_accept_more_false(context: Any) -> None:
    """Verify can_accept_more returns false."""
    assert context.scheduler.can_accept_more() is False


@then("{count:d} subplans should complete successfully")
def step_verify_count_succeeded(context: Any, count: int) -> None:
    """Verify count of successful subplans."""
    assert context.result is not None
    successful = [
        s for s in context.result.statuses if s.status == ProcessingState.COMPLETE
    ]
    assert len(successful) == count


@then("{count:d} subplans should be errored")
def step_verify_count_errored(context: Any, count: int) -> None:
    """Verify count of errored subplans."""
    assert context.result is not None
    errored = [
        s for s in context.result.statuses if s.status == ProcessingState.ERRORED
    ]
    assert len(errored) == count


@then("at least one subplan should be errored with timeout")
def step_verify_timeout_error(context: Any) -> None:
    """Verify timeout error."""
    assert context.result is not None
    for status in context.result.statuses:
        if status.status == ProcessingState.ERRORED:
            assert "timeout" in (status.error or "").lower()
            return
    raise AssertionError("No timeout error found")


@then("the second subplan should complete successfully")
def step_verify_second_succeeds(context: Any) -> None:
    """Verify second subplan succeeded.

    Looks up by subplan_id rather than by index because parallel execution
    may return statuses in completion order.
    """
    assert context.result is not None
    second_id = _to_ulid_id("subplan-001")
    status_map = {s.subplan_id: s for s in context.result.statuses}
    assert status_map[second_id].status == ProcessingState.COMPLETE


@then("subplan A should complete before subplan B")
def step_verify_a_before_b(context: Any) -> None:
    """Verify A completes before B."""
    assert context.result is not None
    status_map = {s.subplan_id: s for s in context.result.statuses}
    a_id = _to_ulid_id("subplan-A")
    b_id = _to_ulid_id("subplan-B")
    assert status_map[a_id].completed_at < status_map[b_id].completed_at


@then("subplan B should complete before subplan C")
def step_verify_b_before_c(context: Any) -> None:
    """Verify B completes before C."""
    assert context.result is not None
    status_map = {s.subplan_id: s for s in context.result.statuses}
    b_id = _to_ulid_id("subplan-B")
    c_id = _to_ulid_id("subplan-C")
    assert status_map[b_id].completed_at < status_map[c_id].completed_at


@then("the dependency order should be respected")
def step_verify_dependency_order(context: Any) -> None:
    """Verify dependency order is respected."""
    assert context.result is not None
    status_map = {s.subplan_id: s for s in context.result.statuses}
    a_id = _to_ulid_id("subplan-A")
    b_id = _to_ulid_id("subplan-B")
    c_id = _to_ulid_id("subplan-C")
    d_id = _to_ulid_id("subplan-D")

    assert status_map[a_id].completed_at < status_map[b_id].completed_at
    assert status_map[a_id].completed_at < status_map[c_id].completed_at

    assert status_map[b_id].completed_at < status_map[d_id].completed_at
    assert status_map[c_id].completed_at < status_map[d_id].completed_at


@then("the peak concurrent execution should be at least {min_concurrent:d}")
def step_verify_min_concurrency(context: Any, min_concurrent: int) -> None:
    """Verify minimum concurrency."""
    if not hasattr(context, "execution_start_times"):
        return  # Skip if not tracking

    # Calculate peak concurrency
    events = []
    for subplan_id, start_time in context.execution_start_times.items():
        end_time = context.execution_end_times.get(subplan_id, start_time)
        events.append((start_time, "start"))
        events.append((end_time, "end"))

    events.sort()
    current_concurrent = 0
    peak_concurrent = 0

    for _, event_type in events:
        if event_type == "start":
            current_concurrent += 1
            peak_concurrent = max(peak_concurrent, current_concurrent)
        else:
            current_concurrent -= 1

    assert peak_concurrent >= min_concurrent


@then("the first subplan should have {count:d} previous attempt recorded")
def step_verify_previous_attempts(context: Any, count: int) -> None:
    """Verify previous attempts."""
    assert context.result is not None
    assert len(context.result.statuses[0].previous_attempts) == count


@then("both subplans should complete successfully")
def step_verify_both_succeed(context: Any) -> None:
    """Verify both subplans succeeded."""
    assert context.result is not None
    assert len(context.result.statuses) == 2
    for status in context.result.statuses:
        assert status.status == ProcessingState.COMPLETE


# --- Additional Given steps ---


@given("{count:d} subplans to execute with staggered completion")
def step_create_subplans_to_execute_staggered(context: Any, count: int) -> None:
    """Create subplans with staggered completion."""
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    block_seconds = {_to_ulid_id(f"subplan-{i:03d}"): i * 0.05 for i in range(count)}
    context.scheduler._executor_fn = create_executor_fn(
        context, block_seconds=block_seconds
    )


@given(
    "{count:d} subplans where the first will fail once with TimeoutError then succeed"
)
def step_create_subplans_retry_then_succeed(context: Any, count: int) -> None:
    """Create subplans where first fails once then succeeds (retry scenario)."""
    context.subplans = [create_subplan_status(f"subplan-{i:03d}") for i in range(count)]
    first_id = _to_ulid_id("subplan-000")
    # Stateful executor: fails first call for first_id, succeeds thereafter.
    fail_state = {"called_for_first": False}

    def executor(status: SubplanStatus) -> SubplanExecutionOutput:
        if status.subplan_id == first_id and not fail_state["called_for_first"]:
            fail_state["called_for_first"] = True
            return SubplanExecutionOutput(
                subplan_id=status.subplan_id,
                success=False,
                error="TimeoutError: simulated timeout",
            )
        return SubplanExecutionOutput(
            subplan_id=status.subplan_id,
            success=True,
            files={"test.txt": "content"},
            files_changed=1,
        )

    context.scheduler._executor_fn = executor


@given(
    "a parallel subplan scheduler with max_parallel {max_parallel:d} with fail_fast disabled"
)
def step_create_scheduler_with_fail_fast_disabled(
    context: Any, max_parallel: int
) -> None:
    """Create a scheduler with fail_fast disabled."""
    config = SubplanConfig(
        execution_mode=ExecutionMode.PARALLEL,
        max_parallel=max_parallel,
        fail_fast=False,
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


@given("a parallel subplan scheduler in {mode} mode")
def step_create_scheduler_mode_only(context: Any, mode: str) -> None:
    """Create a scheduler with specified execution mode (default max_parallel)."""
    mode_map = {
        "SEQUENTIAL": ExecutionMode.SEQUENTIAL,
        "PARALLEL": ExecutionMode.PARALLEL,
        "DEPENDENCY_ORDERED": ExecutionMode.DEPENDENCY_ORDERED,
    }
    config = SubplanConfig(
        execution_mode=mode_map.get(mode, ExecutionMode.PARALLEL),
    )
    executor_fn = create_executor_fn(context)
    context.scheduler = ParallelSubplanScheduler(
        config=config,
        executor_fn=executor_fn,
    )


# --- Additional Then steps ---
# Most error-style assertions (`a config validation error should be raised`,
# `the first subplan should complete successfully`, `the execution result
# should report all succeeded`, etc.) are shared with subplan_execution_steps.py
# via behave's global step registry — they live there, not here. Only steps
# unique to the parallel_subplan_scheduler scenarios are defined below.


@then("the first subplan should be errored with timeout")
def step_verify_first_errored_with_timeout(context: Any) -> None:
    """Verify first subplan errored with timeout.

    Looks up by subplan_id rather than by index because parallel execution
    may return statuses in completion order.
    """
    assert context.result is not None
    first_id = _to_ulid_id("subplan-000")
    status_map = {s.subplan_id: s for s in context.result.statuses}
    first = status_map[first_id]
    assert first.status == ProcessingState.ERRORED
    assert "timeout" in (first.error or "").lower()