"""Step definitions for features/acms_pipeline_orchestrator.feature.

Tests the ACMS Pipeline Orchestrator and Phase 1 production components:
ConfidenceWeightedSelector, ProportionalBudgetAllocator,
ParallelStrategyExecutor, CircuitBreaker, ContextAssemblyPipeline,
and StageTimings.
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

from behave import given, then, when
from behave.runner import Context

from cleveragents.application.services.acms_pipeline import (
    CircuitBreaker,
    ConfidenceWeightedSelector,
    ContextAssemblyPipeline,
    ParallelStrategyExecutor,
    ProportionalBudgetAllocator,
    StageTimings,
)
from cleveragents.application.services.acms_service import (
    DefaultStrategySelector,
    RecencyStrategy,
    RelevanceStrategy,
    StrategyCapabilities,
    TieredStrategy,
)
from cleveragents.application.services.acms_skeleton_compressor import (
    DepthReductionCompressor,
)
from cleveragents.domain.models.core.context_fragment import (
    ContextBudget,
    ContextFragment,
    FragmentProvenance,
)

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

_DEFAULT_PROVENANCE = FragmentProvenance(resource_uri="test://orchestrator")


def _make_frag(**kwargs: Any) -> ContextFragment:
    """Create a ContextFragment with test defaults."""
    kwargs.setdefault("uko_node", "test://orchestrator")
    kwargs.setdefault("token_count", 0)
    kwargs.setdefault("provenance", _DEFAULT_PROVENANCE)
    return ContextFragment(**kwargs)


class _ZeroConfidenceStrategy:
    """Test strategy that always returns 0.0 confidence."""

    @property
    def name(self) -> str:
        return "zero_conf"

    @property
    def capabilities(self) -> StrategyCapabilities:
        return StrategyCapabilities()

    def can_handle(self, request: dict[str, Any]) -> float:
        return 0.0

    def assemble(
        self,
        fragments: Sequence[ContextFragment],
        budget: ContextBudget,
    ) -> Sequence[ContextFragment]:
        return []

    def explain(self) -> str:
        return "Always zero confidence."


class _TrackingTestStrategy:
    """Test strategy that tracks invocation and returns all fragments within budget."""

    def __init__(self, name: str = "tracking") -> None:
        self._name = name
        self.invoked = False

    @property
    def name(self) -> str:
        return self._name

    @property
    def capabilities(self) -> StrategyCapabilities:
        return StrategyCapabilities()

    def can_handle(self, request: dict[str, Any]) -> float:
        return 0.7

    def assemble(
        self,
        fragments: Sequence[ContextFragment],
        budget: ContextBudget,
    ) -> Sequence[ContextFragment]:
        self.invoked = True
        result: list[ContextFragment] = []
        total = 0
        for frag in fragments:
            if total + frag.token_count <= budget.available_tokens:
                result.append(frag)
                total += frag.token_count
        return result

    def explain(self) -> str:
        return "Tracking strategy for testing."


class _FailingStrategy:
    """Test strategy that always raises an exception."""

    @property
    def name(self) -> str:
        return "failing"

    @property
    def capabilities(self) -> StrategyCapabilities:
        return StrategyCapabilities()

    def can_handle(self, request: dict[str, Any]) -> float:
        return 0.5

    def assemble(
        self,
        fragments: Sequence[ContextFragment],
        budget: ContextBudget,
    ) -> Sequence[ContextFragment]:
        msg = "Intentional test failure"
        raise RuntimeError(msg)

    def explain(self) -> str:
        return "Always fails."


class _ReverseOrchestratorStrategy:
    """Test strategy that reverses fragment order."""

    @property
    def name(self) -> str:
        return "reverse_orch"

    @property
    def capabilities(self) -> StrategyCapabilities:
        return StrategyCapabilities()

    def can_handle(self, request: dict[str, Any]) -> float:
        return 0.5

    def assemble(
        self,
        fragments: Sequence[ContextFragment],
        budget: ContextBudget,
    ) -> Sequence[ContextFragment]:
        reversed_frags = list(reversed(fragments))
        result: list[ContextFragment] = []
        total = 0
        for frag in reversed_frags:
            if total + frag.token_count <= budget.available_tokens:
                result.append(frag)
                total += frag.token_count
        return result

    def explain(self) -> str:
        return "Reverses fragment order."


class _TrackingOrchestratorSelector:
    """Strategy selector that tracks invocation and delegates to default."""

    def __init__(self) -> None:
        self.called = False

    def select(
        self,
        strategies: Sequence[Any],
        request: dict[str, Any],
    ) -> list[tuple[Any, float]]:
        self.called = True
        return DefaultStrategySelector().select(strategies, request)


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


@given("the pipeline orchestrator modules are available")
def step_orchestrator_modules_available(context: Context) -> None:
    """Ensure all pipeline orchestrator modules are importable."""
    pass


@given("a test strategy with zero confidence")
def step_zero_confidence_strategy(context: Context) -> None:
    context.zero_conf_strategy = _ZeroConfidenceStrategy()


@given("a circuit breaker with threshold {threshold:d}")
def step_circuit_breaker(context: Context, threshold: int) -> None:
    context.circuit_breaker = CircuitBreaker(failure_threshold=threshold)


@given("a tracking test strategy")
def step_tracking_strategy(context: Context) -> None:
    context.tracking_strategy = _TrackingTestStrategy()


@given("context fragments for executor testing")
def step_executor_fragments(context: Context) -> None:
    context.executor_fragments = [
        _make_frag(
            uko_node="project://test/a.py",
            content="alpha",
            token_count=50,
            relevance_score=0.9,
        ),
        _make_frag(
            uko_node="project://test/b.py",
            content="beta",
            token_count=50,
            relevance_score=0.5,
        ),
    ]


@given('a ParallelStrategyExecutor with a pre-broken circuit for "{name}"')
def step_pre_broken_executor(context: Context, name: str) -> None:
    cb = CircuitBreaker(failure_threshold=1)
    cb.record_failure(name)
    context.broken_executor = ParallelStrategyExecutor(circuit_breaker=cb)
    context.broken_strategy_name = name


@given("a strategy that always raises an exception")
def step_failing_strategy(context: Context) -> None:
    context.failing_strategy = _FailingStrategy()


@given("a ContextAssemblyPipeline with default components")
def step_default_pipeline(context: Context) -> None:
    context.orch_pipeline = ContextAssemblyPipeline()


@then("the orchestrator pipeline should use DepthReductionCompressor")
def step_pipeline_uses_depth_reduction(context: Context) -> None:
    compressor = getattr(context.orch_pipeline, "_skeleton_compressor", None)
    assert isinstance(compressor, DepthReductionCompressor), (
        "Expected ContextAssemblyPipeline to default to DepthReductionCompressor, "
        f"got {type(compressor).__name__ if compressor is not None else None}"
    )


@given("the following orchestrator test fragments:")
def step_orch_fragments_table(context: Context) -> None:
    context.orch_fragments = []
    for row in context.table:
        uko_node = row["uko_node"]
        frag = ContextFragment(
            uko_node=uko_node,
            content=row["content"],
            relevance_score=float(row["score"]),
            token_count=int(row["tokens"]),
            provenance=FragmentProvenance(resource_uri=uko_node),
        )
        context.orch_fragments.append(frag)


@given("a pipeline budget with max_tokens {max_tok:d} and reserved_tokens {res:d}")
def step_pipeline_budget(context: Context, max_tok: int, res: int) -> None:
    context.orch_budget = ContextBudget(max_tokens=max_tok, reserved_tokens=res)


@given("a ContextAssemblyPipeline with a custom tracking selector")
def step_pipeline_custom_selector(context: Context) -> None:
    selector = _TrackingOrchestratorSelector()
    context.orch_custom_selector = selector
    context.orch_pipeline = ContextAssemblyPipeline(strategy_selector=selector)


@given("a custom orchestrator strategy that returns fragments in reverse")
def step_custom_reverse_strategy(context: Context) -> None:
    if not hasattr(context, "orch_pipeline"):
        context.orch_pipeline = ContextAssemblyPipeline()
    context.orch_pipeline.register_strategy(
        "reverse_orch", _ReverseOrchestratorStrategy()
    )


# ---------------------------------------------------------------------------
# When steps — ConfidenceWeightedSelector
# ---------------------------------------------------------------------------


@when("I use ConfidenceWeightedSelector to select from 3 strategies")
def step_cws_select(context: Context) -> None:
    strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
    selector = ConfidenceWeightedSelector()
    context.cws_results = selector.select(strategies, {})


@when('I select with preferred_strategies containing "{preferred}"')
def step_cws_preferred(context: Context, preferred: str) -> None:
    strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
    selector = ConfidenceWeightedSelector(preference_boost=1.5)
    context.cws_results_preferred = selector.select(
        strategies, {"preferred_strategies": [preferred]}
    )
    # Also get non-boosted results for comparison
    context.cws_results_baseline = selector.select(strategies, {})
    context.preferred_name = preferred


@when("I use ConfidenceWeightedSelector with the zero-confidence strategy")
def step_cws_zero_conf(context: Context) -> None:
    strategies = [RelevanceStrategy(), context.zero_conf_strategy]
    selector = ConfidenceWeightedSelector()
    context.cws_results_with_zero = selector.select(strategies, {})


# ---------------------------------------------------------------------------
# When steps — ProportionalBudgetAllocator
# ---------------------------------------------------------------------------


@when(
    "I use ProportionalBudgetAllocator for candidates with confidences {c1:g} and {c2:g} and budget {budget:d}"
)
def step_pba_proportional(context: Context, c1: float, c2: float, budget: int) -> None:
    s1 = RelevanceStrategy()
    s2 = RecencyStrategy()
    allocator = ProportionalBudgetAllocator()
    context.pba_allocations = allocator.allocate([(s1, c1), (s2, c2)], budget)
    context.pba_budget = budget


@when(
    "I use ProportionalBudgetAllocator with min_useful_budget {min_budget:d} for candidates with confidences {c1:g} and {c2:g} and budget {budget:d}"
)
def step_pba_min_budget(
    context: Context, min_budget: int, c1: float, c2: float, budget: int
) -> None:
    s1 = RelevanceStrategy()
    s2 = RecencyStrategy()
    allocator = ProportionalBudgetAllocator(min_useful_budget=min_budget)
    context.pba_allocations = allocator.allocate([(s1, c1), (s2, c2)], budget)
    context.pba_budget = budget


@when("I use ProportionalBudgetAllocator for a single candidate with budget {budget:d}")
def step_pba_single(context: Context, budget: int) -> None:
    s1 = RelevanceStrategy()
    allocator = ProportionalBudgetAllocator()
    context.pba_allocations = allocator.allocate([(s1, 0.8)], budget)
    context.pba_budget = budget


@when(
    "I use ProportionalBudgetAllocator for {count:d} zero-confidence candidates with budget {budget:d}"
)
def step_pba_zero_conf(context: Context, count: int, budget: int) -> None:
    strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
    allocator = ProportionalBudgetAllocator()
    candidates = [(s, 0.0) for s in strategies[:count]]
    context.pba_allocations = allocator.allocate(candidates, budget)
    context.pba_budget = budget


@when(
    "I use ProportionalBudgetAllocator with min_useful_budget {min_budget:d} for {count:d} candidates with budget {budget:d}"
)
def step_pba_all_excluded(
    context: Context, min_budget: int, count: int, budget: int
) -> None:
    strategies = [RelevanceStrategy(), RecencyStrategy()]
    allocator = ProportionalBudgetAllocator(min_useful_budget=min_budget)
    candidates = [(s, 0.5) for s in strategies[:count]]
    context.pba_allocations = allocator.allocate(candidates, budget)
    context.pba_budget = budget


# ---------------------------------------------------------------------------
# When steps — CircuitBreaker
# ---------------------------------------------------------------------------


@when('I record {count:d} consecutive failures for strategy "{name}"')
def step_cb_failures(context: Context, count: int, name: str) -> None:
    for _ in range(count):
        context.circuit_breaker.record_failure(name)


@when('I record 2 failures then 1 success for strategy "{name}"')
def step_cb_failures_then_success(context: Context, name: str) -> None:
    context.circuit_breaker.record_failure(name)
    context.circuit_breaker.record_failure(name)
    context.circuit_breaker.record_success(name)


@when('I record 2 failures for strategy "{name}"')
def step_cb_2_failures(context: Context, name: str) -> None:
    context.circuit_breaker.record_failure(name)
    context.circuit_breaker.record_failure(name)


@when('I reset the circuit for "{name}"')
def step_cb_reset(context: Context, name: str) -> None:
    context.circuit_breaker.reset(name)


@when('I record 1 failure for strategy "{name1}" and 1 failure for strategy "{name2}"')
def step_cb_two_strat_failures(context: Context, name1: str, name2: str) -> None:
    context.circuit_breaker.record_failure(name1)
    context.circuit_breaker.record_failure(name2)


@when("I reset all circuits")
def step_cb_reset_all(context: Context) -> None:
    context.circuit_breaker.reset_all()


# ---------------------------------------------------------------------------
# When steps — ParallelStrategyExecutor
# ---------------------------------------------------------------------------


@when("I execute the strategy via ParallelStrategyExecutor")
def step_executor_run(context: Context) -> None:
    executor = ParallelStrategyExecutor()
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    allocations = [(context.tracking_strategy, 0.7, 200)]
    context.executor_results = executor.execute(
        allocations, context.executor_fragments, budget
    )


@when("I execute with the circuit-broken strategy")
def step_executor_broken(context: Context) -> None:
    strategy = _TrackingTestStrategy(name=context.broken_strategy_name)
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    allocations = [(strategy, 0.7, 200)]
    context.executor_results = context.broken_executor.execute(
        allocations, context.executor_fragments, budget
    )


@when("I execute the failing strategy via ParallelStrategyExecutor")
def step_executor_failing(context: Context) -> None:
    cb = CircuitBreaker()
    executor = ParallelStrategyExecutor(circuit_breaker=cb)
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    allocations = [(context.failing_strategy, 0.5, 200)]
    context.executor_results = executor.execute(
        allocations, context.executor_fragments, budget
    )
    context.executor_cb = cb


@when("I execute with zero-budget allocation")
def step_executor_zero_budget(context: Context) -> None:
    executor = ParallelStrategyExecutor()
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    allocations = [(context.tracking_strategy, 0.7, 0)]
    context.executor_results = executor.execute(
        allocations, context.executor_fragments, budget
    )


@when("I execute 2 tracking strategies in parallel via ParallelStrategyExecutor")
def step_executor_parallel(context: Context) -> None:
    s1 = _TrackingTestStrategy(name="parallel_a")
    s2 = _TrackingTestStrategy(name="parallel_b")
    context.parallel_strategies = [s1, s2]
    executor = ParallelStrategyExecutor(max_workers=2)
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    allocations: list[tuple[Any, float, int]] = [
        (s1, 0.6, 100),
        (s2, 0.4, 100),
    ]
    context.executor_results = executor.execute(
        allocations, context.executor_fragments, budget
    )


@when("I execute a tracking strategy and a failing strategy in parallel")
def step_executor_parallel_mixed(context: Context) -> None:
    s_ok = _TrackingTestStrategy(name="good_parallel")
    s_fail = _FailingStrategy()
    context.parallel_ok_strategy = s_ok
    cb = CircuitBreaker()
    executor = ParallelStrategyExecutor(max_workers=2, circuit_breaker=cb)
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    allocations: list[tuple[Any, float, int]] = [
        (s_ok, 0.6, 100),
        (s_fail, 0.4, 100),
    ]
    context.executor_results = executor.execute(
        allocations, context.executor_fragments, budget
    )
    context.parallel_cb = cb


@when("I execute with empty allocations via ParallelStrategyExecutor")
def step_executor_empty_allocations(context: Context) -> None:
    executor = ParallelStrategyExecutor()
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    context.executor_results = executor.execute([], context.executor_fragments, budget)


@when("I create a ParallelStrategyExecutor with a custom circuit breaker")
def step_executor_custom_cb(context: Context) -> None:
    custom_cb = CircuitBreaker(failure_threshold=10)
    context.custom_executor = ParallelStrategyExecutor(circuit_breaker=custom_cb)
    context.custom_cb = custom_cb


@when("I use ProportionalBudgetAllocator for empty candidates with budget {budget:d}")
def step_pba_empty_candidates(context: Context, budget: int) -> None:
    allocator = ProportionalBudgetAllocator()
    context.pba_allocations = allocator.allocate([], budget)
    context.pba_budget = budget


# ---------------------------------------------------------------------------
# When steps — ContextAssemblyPipeline
# ---------------------------------------------------------------------------


@when('I assemble via the orchestrator with strategy "{strategy}"')
def step_orch_assemble(context: Context, strategy: str) -> None:
    context.orch_error = None
    try:
        context.orch_payload = context.orch_pipeline.assemble(
            plan_id="01JQTESTPN00000000000000AA",
            fragments=list(context.orch_fragments),
            budget=context.orch_budget,
            strategy=strategy,
        )
    except ValueError as exc:
        context.orch_error = exc


@when('I assemble via the orchestrator with invalid plan_id "{plan_id}"')
def step_orch_assemble_invalid(context: Context, plan_id: str) -> None:
    context.orch_error = None
    try:
        context.orch_payload = context.orch_pipeline.assemble(
            plan_id=plan_id,
            fragments=list(context.orch_fragments),
            budget=context.orch_budget,
        )
    except ValueError as exc:
        context.orch_error = exc


# ---------------------------------------------------------------------------
# When steps — StageTimings
# ---------------------------------------------------------------------------


@when("I create a StageTimings with total_ms {total:g}")
def step_create_timings(context: Context, total: float) -> None:
    context.test_timings = StageTimings(total_ms=total)


# ---------------------------------------------------------------------------
# Then steps — ConfidenceWeightedSelector
# ---------------------------------------------------------------------------


@then("all strategies with positive confidence should be returned")
def step_cws_all_positive(context: Context) -> None:
    assert len(context.cws_results) == 3, (
        f"Expected 3 strategies, got {len(context.cws_results)}"
    )
    for _, conf in context.cws_results:
        assert conf > 0.0


@then("the results should be sorted by confidence descending")
def step_cws_sorted(context: Context) -> None:
    confs = [c for _, c in context.cws_results]
    assert confs == sorted(confs, reverse=True), (
        f"Results not sorted descending: {confs}"
    )


@then('the "{name}" strategy should have a boosted confidence')
def step_cws_boosted(context: Context, name: str) -> None:
    boosted = {s.name: c for s, c in context.cws_results_preferred}
    baseline = {s.name: c for s, c in context.cws_results_baseline}
    assert boosted[name] >= baseline[name], (
        f"Expected {name} boosted ({boosted[name]}) >= baseline ({baseline[name]})"
    )


@then("the boosted confidence should not exceed 1.0")
def step_cws_boosted_capped(context: Context) -> None:
    for _, conf in context.cws_results_preferred:
        assert conf <= 1.0, f"Confidence {conf} exceeds 1.0"


@then("the zero-confidence strategy should not be in the results")
def step_cws_no_zero(context: Context) -> None:
    names = [s.name for s, _ in context.cws_results_with_zero]
    assert "zero_conf" not in names, (
        f"Zero-confidence strategy should be excluded, got {names}"
    )


# ---------------------------------------------------------------------------
# Then steps — ProportionalBudgetAllocator
# ---------------------------------------------------------------------------


@then(
    "the first candidate should receive approximately {tokens:d} tokens from proportional allocator"
)
def step_pba_first(context: Context, tokens: int) -> None:
    actual = context.pba_allocations[0][2]
    assert abs(actual - tokens) <= 1, (
        f"First candidate got {actual}, expected ~{tokens}"
    )


@then(
    "the second candidate should receive approximately {tokens:d} tokens from proportional allocator"
)
def step_pba_second(context: Context, tokens: int) -> None:
    actual = context.pba_allocations[1][2]
    assert abs(actual - tokens) <= 1, (
        f"Second candidate got {actual}, expected ~{tokens}"
    )


@then("the total proportional allocation should equal exactly {budget:d}")
def step_pba_total_exact(context: Context, budget: int) -> None:
    total = sum(a[2] for a in context.pba_allocations)
    assert total == budget, f"Total allocation {total} != {budget}"


@then("only the high-confidence candidate should receive tokens")
def step_pba_only_high(context: Context) -> None:
    assert len(context.pba_allocations) == 1, (
        f"Expected 1 allocation, got {len(context.pba_allocations)}"
    )


@then("the single candidate should receive all {budget:d} tokens")
def step_pba_single_all(context: Context, budget: int) -> None:
    assert len(context.pba_allocations) == 1
    assert context.pba_allocations[0][2] == budget


@then("each proportional allocation should be {low:d} or {high:d}")
def step_pba_each_bounded(context: Context, low: int, high: int) -> None:
    for _, _, tokens in context.pba_allocations:
        assert tokens in {low, high}, (
            f"Expected allocation to be {low} or {high}, got {tokens}"
        )


@then("the highest-confidence candidate should receive the full budget")
def step_pba_fallback(context: Context) -> None:
    # When all candidates are excluded, the allocator falls back to the
    # highest-confidence candidate with the full budget.
    total = sum(a[2] for a in context.pba_allocations)
    assert total == context.pba_budget, (
        f"Expected full budget {context.pba_budget}, got {total}"
    )


# ---------------------------------------------------------------------------
# Then steps — CircuitBreaker
# ---------------------------------------------------------------------------


@then('the circuit for "{name}" should be open')
def step_cb_is_open(context: Context, name: str) -> None:
    assert context.circuit_breaker.is_open(name), (
        f"Expected circuit for '{name}' to be open"
    )


@then('the circuit for "{name}" should be closed')
def step_cb_is_closed(context: Context, name: str) -> None:
    assert not context.circuit_breaker.is_open(name), (
        f"Expected circuit for '{name}' to be closed"
    )


# ---------------------------------------------------------------------------
# Then steps — ParallelStrategyExecutor
# ---------------------------------------------------------------------------


@then("the executor should return fragments from the strategy")
def step_executor_has_results(context: Context) -> None:
    assert len(context.executor_results) > 0, "Expected fragments from executor"


@then("the tracking strategy should have been invoked")
def step_tracking_invoked(context: Context) -> None:
    assert context.tracking_strategy.invoked, "Tracking strategy was not invoked"


@then("the executor should return {count:d} fragments")
def step_executor_count(context: Context, count: int) -> None:
    assert len(context.executor_results) == count, (
        f"Expected {count} fragments, got {len(context.executor_results)}"
    )


@then('the circuit breaker should have recorded a failure for "{name}"')
def step_cb_has_failure(context: Context, name: str) -> None:
    # After one failure, the circuit isn't necessarily open (depends on threshold)
    # but the failure count should be > 0
    assert context.executor_cb._failures.get(name, 0) > 0, (
        f"Expected failure recorded for '{name}'"
    )


@then("the executor should return fragments from all parallel strategies")
def step_executor_parallel_results(context: Context) -> None:
    # Two strategies each returning up to their budget from the same 2 fragments
    assert len(context.executor_results) > 0, (
        "Expected fragments from parallel execution"
    )


@then("both parallel strategies should have been invoked")
def step_parallel_both_invoked(context: Context) -> None:
    for s in context.parallel_strategies:
        assert s.invoked, f"Strategy {s.name} was not invoked"


@then("the executor should return fragments only from the successful parallel strategy")
def step_executor_parallel_mixed_results(context: Context) -> None:
    assert len(context.executor_results) > 0, (
        "Expected fragments from the successful strategy"
    )
    assert context.parallel_ok_strategy.invoked, (
        "Good strategy should have been invoked"
    )


@then('the parallel circuit breaker should record a failure for "{name}"')
def step_parallel_cb_failure(context: Context, name: str) -> None:
    assert context.parallel_cb._failures.get(name, 0) > 0, (
        f"Expected failure recorded for '{name}'"
    )


@then("the circuit_breaker property should return the custom breaker")
def step_executor_cb_property(context: Context) -> None:
    assert context.custom_executor.circuit_breaker is context.custom_cb, (
        "circuit_breaker property did not return the custom breaker"
    )


@then("the proportional allocator should return an empty list")
def step_pba_empty_result(context: Context) -> None:
    assert len(context.pba_allocations) == 0, (
        f"Expected empty allocations, got {len(context.pba_allocations)}"
    )


# ---------------------------------------------------------------------------
# Then steps — ContextAssemblyPipeline
# ---------------------------------------------------------------------------


@then("the orchestrator payload should contain {count:d} fragments")
def step_orch_payload_count(context: Context, count: int) -> None:
    assert len(context.orch_payload.fragments) == count, (
        f"Expected {count} fragments, got {len(context.orch_payload.fragments)}"
    )


@then("the orchestrator should have timing metadata")
def step_orch_has_timings(context: Context) -> None:
    timings = context.orch_pipeline.last_timings
    assert timings is not None, "Expected timing metadata to be populated"
    context.orch_timings = timings


@then("all timing values should be non-negative")
def step_orch_timings_nonneg(context: Context) -> None:
    timings = context.orch_pipeline.last_timings
    assert timings is not None
    for name in type(timings).model_fields:
        val = getattr(timings, name)
        assert val >= 0.0, f"Timing {name} = {val} is negative"


@then('the orchestrator payload strategies used should include "{strategy}"')
def step_orch_strategies(context: Context, strategy: str) -> None:
    assert strategy in context.orch_payload.strategies_used, (
        f"Expected {strategy!r} in {context.orch_payload.strategies_used}"
    )


@then("the orchestrator payload should be within budget")
def step_orch_within_budget(context: Context) -> None:
    assert context.orch_payload.is_within_budget


@then('an orchestrator ValueError should be raised mentioning "{keyword}"')
def step_orch_error(context: Context, keyword: str) -> None:
    assert context.orch_error is not None, "Expected ValueError but none raised"
    assert keyword.lower() in str(context.orch_error).lower(), (
        f"Expected '{keyword}' in error: {context.orch_error}"
    )


@then("the custom orchestrator selector should have been called")
def step_orch_selector_called(context: Context) -> None:
    assert context.orch_custom_selector.called, (
        "Custom orchestrator selector was not invoked"
    )


@then('the first orchestrator fragment content should be "{content}"')
def step_orch_first_content(context: Context, content: str) -> None:
    assert context.orch_payload.fragments[0].content == content, (
        f"Expected first fragment '{content}', "
        f"got '{context.orch_payload.fragments[0].content}'"
    )


# ---------------------------------------------------------------------------
# Then steps — StageTimings
# ---------------------------------------------------------------------------


@then("the timings total_ms should be {total:g}")
def step_timings_total(context: Context, total: float) -> None:
    assert context.test_timings.total_ms == total


@then("the timings should have all 10 named fields")
def step_timings_fields(context: Context) -> None:
    field_names = set(StageTimings.model_fields.keys())
    assert len(field_names) == 10, f"Expected 10 fields, got {len(field_names)}"
    expected_names = {
        "strategy_selection_ms",
        "budget_allocation_ms",
        "strategy_execution_ms",
        "deduplication_ms",
        "depth_resolution_ms",
        "scoring_ms",
        "packing_ms",
        "ordering_ms",
        "preamble_generation_ms",
        "total_ms",
    }
    assert field_names == expected_names, (
        f"Field mismatch: {field_names.symmetric_difference(expected_names)}"
    )
