cleveragents-core/features/steps/acms_pipeline_steps.py

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

Tests the ACMS v1 context assembly pipeline domain models and service
directly in-memory -- no database required.
"""

from __future__ import annotations

from collections.abc import Sequence
from datetime import UTC, datetime, timedelta
from typing import Any

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

from cleveragents.application.services.acms_service import (
    ACMSPipeline,
    DefaultBudgetAllocator,
    DefaultSkeletonCompressor,
    DefaultStrategySelector,
    RecencyStrategy,
    RelevanceStrategy,
    StrategyCapabilities,
    TieredStrategy,
)
from cleveragents.domain.models.acms.crp import ContextRequest
from cleveragents.domain.models.core.context_fragment import (
    ContextBudget,
    ContextFragment,
    ContextPayload,
    FragmentProvenance,
)

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

# Default provenance for test fragments.
_DEFAULT_PROVENANCE = FragmentProvenance(resource_uri="test://default")


def _make_fragment(**kwargs: Any) -> ContextFragment:
    """Create a ContextFragment with sensible test defaults.

    Callers can override any field.  ``uko_node``, ``token_count``, and
    ``provenance`` get defaults so that most test steps don't need to
    specify them explicitly.
    """
    kwargs.setdefault("uko_node", "test://default")
    kwargs.setdefault("token_count", 0)
    kwargs.setdefault("provenance", _DEFAULT_PROVENANCE)
    return ContextFragment(**kwargs)


class _ReverseContextStrategy:
    """Test strategy that reverses fragment order and fits to budget."""

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

    @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
        available = budget.available_tokens
        for frag in reversed_frags:
            if total >= available:
                break
            if total + frag.token_count <= available:
                result.append(frag)
                total += frag.token_count
        return result

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


# ---------------------------------------------------------------------------
# Given — ContextFragment
# ---------------------------------------------------------------------------


@given(
    'a context fragment with uko_node "{uko_node}" and content "{content}" and token_count {tokens:d}'
)
def step_fragment_with_defaults(
    context: Context, uko_node: str, content: str, tokens: int
) -> None:
    context.fragment = ContextFragment(
        uko_node=uko_node,
        content=content,
        token_count=tokens,
        provenance=FragmentProvenance(resource_uri=uko_node),
    )


@given("a context fragment with all fields:")
def step_fragment_with_all_fields(context: Context) -> None:
    data: dict[str, str] = {}
    for row in context.table:
        data[row["field"]] = row["value"]
    meta: dict[str, str] = {}
    if "meta_key" in data and "meta_val" in data:
        meta[data["meta_key"]] = data["meta_val"]
    resource_uri = data.get("resource_uri", data["uko_node"])
    context.fragment = ContextFragment(
        uko_node=data["uko_node"],
        content=data["content"],
        relevance_score=float(data["score"]),
        token_count=int(data["tokens"]),
        detail_depth=int(data.get("detail_depth", "0")),
        tier=data["tier"],
        metadata=meta,
        provenance=FragmentProvenance(resource_uri=resource_uri),
    )


# ---------------------------------------------------------------------------
# Given — ContextBudget
# ---------------------------------------------------------------------------


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


# ---------------------------------------------------------------------------
# Given — Fragment collections
# ---------------------------------------------------------------------------


@given("the following context fragments:")
def step_fragments_table(context: Context) -> None:
    context.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.fragments.append(frag)


@given("context fragments with different timestamps:")
def step_fragments_with_timestamps(context: Context) -> None:
    context.fragments = []
    for row in context.table:
        uko_node = row["uko_node"]
        frag = ContextFragment(
            uko_node=uko_node,
            content=row["content"],
            token_count=int(row["tokens"]),
            created_at=datetime.fromisoformat(row["created_at"]),
            provenance=FragmentProvenance(resource_uri=uko_node),
        )
        context.fragments.append(frag)


@given("the following tiered context fragments:")
def step_tiered_fragments(context: Context) -> None:
    context.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"]),
            tier=row["tier"],
            provenance=FragmentProvenance(resource_uri=uko_node),
        )
        context.fragments.append(frag)


@given("no context fragments")
def step_no_fragments(context: Context) -> None:
    context.fragments = []


@given("a custom context strategy that reverses fragment order")
def step_register_custom_strategy(context: Context) -> None:
    if not hasattr(context, "pipeline"):
        context.pipeline = ACMSPipeline()
    context.pipeline.register_strategy("reverse", _ReverseContextStrategy())


# ---------------------------------------------------------------------------
# When — Validation
# ---------------------------------------------------------------------------


@given("the ACMS pipeline modules are available")
def step_acms_modules_available(context: Context) -> None:
    # Just ensures the imports are loadable; nothing to store.
    pass


@when("I create a budget with reserved_tokens equal to max_tokens")
def step_budget_reserved_equals_max(context: Context) -> None:
    context.acms_error = None
    try:
        ContextBudget(max_tokens=100, reserved_tokens=100)
    except ValidationError as exc:
        context.acms_error = exc


@when("I create a budget with reserved_tokens greater than max_tokens")
def step_budget_reserved_greater_than_max(context: Context) -> None:
    context.acms_error = None
    try:
        ContextBudget(max_tokens=100, reserved_tokens=200)
    except ValidationError as exc:
        context.acms_error = exc


@when("I create a budget with reserved_tokens {res:d} and max_tokens {max_tok:d}")
def step_budget_specific(context: Context, res: int, max_tok: int) -> None:
    context.acms_error = None
    try:
        ContextBudget(max_tokens=max_tok, reserved_tokens=res)
    except ValidationError as exc:
        context.acms_error = exc


@when('I create a fragment with invalid tier "{tier}"')
def step_invalid_tier(context: Context, tier: str) -> None:
    context.acms_error = None
    try:
        _make_fragment(content="test", tier=tier, token_count=1)
    except ValidationError as exc:
        context.acms_error = exc


@when("I create a fragment with relevance score {score}")
def step_create_fragment_bad_score(context: Context, score: str) -> None:
    context.acms_error = None
    try:
        _make_fragment(content="test", relevance_score=float(score), token_count=1)
    except ValidationError as exc:
        context.acms_error = exc


@when("I create a fragment with token count {tokens:d}")
def step_create_fragment_bad_tokens(context: Context, tokens: int) -> None:
    context.acms_error = None
    try:
        _make_fragment(content="test", token_count=tokens)
    except ValidationError as exc:
        context.acms_error = exc


@when("I create a fragment with empty uko_node")
def step_create_fragment_empty_uko_node(context: Context) -> None:
    context.acms_error = None
    try:
        ContextFragment(
            uko_node="",
            content="test",
            token_count=1,
            provenance=_DEFAULT_PROVENANCE,
        )
    except ValidationError as exc:
        context.acms_error = exc


@when("I create a fragment with detail depth {depth:d}")
def step_create_fragment_bad_depth(context: Context, depth: int) -> None:
    context.acms_error = None
    try:
        _make_fragment(content="test", detail_depth=depth, token_count=1)
    except ValidationError as exc:
        context.acms_error = exc


# ---------------------------------------------------------------------------
# When — Assembly
# ---------------------------------------------------------------------------


@when('I assemble with strategy "{strategy}"')
def step_assemble_with_strategy(context: Context, strategy: str) -> None:
    if not hasattr(context, "pipeline"):
        context.pipeline = ACMSPipeline()
    context.assemble_error = None
    try:
        context.payload = context.pipeline.assemble(
            plan_id="01JQTESTPN00000000000000AA",
            fragments=list(context.fragments),
            budget=context.budget,
            strategy=strategy,
        )
    except ValueError as exc:
        context.assemble_error = exc


@when("I assemble without specifying a strategy")
def step_assemble_default_strategy(context: Context) -> None:
    if not hasattr(context, "pipeline"):
        context.pipeline = ACMSPipeline()
    context.payload = context.pipeline.assemble(
        plan_id="01JQTESTPN00000000000000AA",
        fragments=list(context.fragments),
        budget=context.budget,
    )


@when('I register the reverse strategy as "{name}"')
def step_register_reverse_as(context: Context, name: str) -> None:
    if not hasattr(context, "pipeline"):
        context.pipeline = ACMSPipeline()
    context.pipeline.register_strategy(name, _ReverseContextStrategy())


# ---------------------------------------------------------------------------
# Then — Fragment assertions
# ---------------------------------------------------------------------------


@then('the fragment content should be "{content}"')
def step_fragment_content(context: Context, content: str) -> None:
    assert context.fragment.content == content


@then("the fragment relevance score should be {score}")
def step_fragment_score(context: Context, score: str) -> None:
    assert context.fragment.relevance_score == float(score)


@then("the fragment token count should be {tokens:d}")
def step_fragment_tokens(context: Context, tokens: int) -> None:
    assert context.fragment.token_count == tokens


@then("the fragment detail depth should be {depth:d}")
def step_fragment_detail_depth(context: Context, depth: int) -> None:
    assert context.fragment.detail_depth == depth


@then("the fragment metadata should be empty")
def step_fragment_metadata_empty(context: Context) -> None:
    assert context.fragment.metadata == {}


@then("the fragment id should be set")
def step_fragment_id_set(context: Context) -> None:
    assert context.fragment.fragment_id
    assert len(context.fragment.fragment_id) > 0


@then("the fragment created_at should be a datetime")
def step_fragment_created_at_is_datetime(context: Context) -> None:
    assert isinstance(context.fragment.created_at, datetime)


@then("the fragment provenance resource_uri should be set")
def step_fragment_provenance_set(context: Context) -> None:
    assert context.fragment.provenance.resource_uri
    assert len(context.fragment.provenance.resource_uri) > 0


@then('the fragment metadata key "{key}" should be "{value}"')
def step_fragment_metadata_key(context: Context, key: str, value: str) -> None:
    assert context.fragment.metadata[key] == value


@then('the fragment should have field "{field_name}"')
def step_fragment_has_field(context: Context, field_name: str) -> None:
    assert hasattr(context.fragment, field_name), (
        f"ContextFragment missing field: {field_name}"
    )


# ---------------------------------------------------------------------------
# Then — Budget assertions
# ---------------------------------------------------------------------------


@then("the available tokens should be {tokens:d}")
def step_available_tokens(context: Context, tokens: int) -> None:
    assert context.budget.available_tokens == tokens


# ---------------------------------------------------------------------------
# Then — Validation assertions
# ---------------------------------------------------------------------------


@then("an ACMS validation error should be raised")
def step_acms_validation_error_raised(context: Context) -> None:
    error = getattr(context, "acms_error", None)
    assert error is not None, "Expected a validation error but none was raised"


@then("an ACMS strategy error should be raised")
def step_acms_strategy_error_raised(context: Context) -> None:
    assert context.assemble_error is not None, (
        "Expected a strategy error but none was raised"
    )


# ---------------------------------------------------------------------------
# Then — Payload assertions
# ---------------------------------------------------------------------------


@then("the payload should contain {count:d} fragments")
def step_payload_fragment_count(context: Context, count: int) -> None:
    assert len(context.payload.fragments) == count


@then('the first fragment content should be "{content}"')
def step_first_fragment_content(context: Context, content: str) -> None:
    assert context.payload.fragments[0].content == content


@then('the second fragment content should be "{content}"')
def step_second_fragment_content(context: Context, content: str) -> None:
    assert context.payload.fragments[1].content == content


@then("the payload total tokens should be {tokens:d}")
def step_payload_total_tokens(context: Context, tokens: int) -> None:
    assert context.payload.total_tokens == tokens


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


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


@then("the payload remaining tokens should be {tokens:d}")
def step_payload_remaining_tokens(context: Context, tokens: int) -> None:
    assert context.payload.remaining_tokens == tokens


@then("the payload budget used should be {fraction}")
def step_payload_budget_used(context: Context, fraction: str) -> None:
    assert abs(context.payload.budget_used - float(fraction)) < 1e-4, (
        f"Expected budget_used={fraction}, got {context.payload.budget_used}"
    )


@then("the payload context hash should be non-empty")
def step_payload_context_hash(context: Context) -> None:
    assert context.payload.context_hash, "Expected non-empty context_hash"
    assert len(context.payload.context_hash) == 64, (
        f"Expected SHA-256 hex (64 chars), got {len(context.payload.context_hash)}"
    )


@then("the payload provenance map should map each fragment ID")
def step_payload_provenance_map(context: Context) -> None:
    for frag in context.payload.fragments:
        assert frag.fragment_id in context.payload.provenance_map, (
            f"Fragment {frag.fragment_id} not in provenance_map"
        )


@then('the payload plan_id should be "{plan_id}"')
def step_payload_plan_id(context: Context, plan_id: str) -> None:
    assert context.payload.plan_id == plan_id, (
        f"Expected plan_id={plan_id!r}, got {context.payload.plan_id!r}"
    )


@then("the payload payload_id should be a non-empty ULID")
def step_payload_payload_id(context: Context) -> None:
    pid = context.payload.payload_id
    assert pid, "payload_id is empty"
    assert len(pid) == 26, f"Expected ULID (26 chars), got {len(pid)} chars: {pid!r}"


@then("the payload assembled_at should be a recent UTC datetime")
def step_payload_assembled_at(context: Context) -> None:
    assembled = context.payload.assembled_at
    assert isinstance(assembled, datetime), (
        f"assembled_at is not datetime, got {type(assembled)}"
    )
    # Should be within the last 60 seconds.
    now = datetime.now(UTC)
    delta = now - assembled
    assert delta < timedelta(seconds=60), (
        f"assembled_at is too old: {assembled} (now={now})"
    )


@then('the payload should have field "{field_name}"')
def step_payload_has_field(context: Context, field_name: str) -> None:
    assert hasattr(context.payload, field_name), (
        f"ContextPayload missing field: {field_name}"
    )


# ---------------------------------------------------------------------------
# T1 — DI injection steps for pipeline components
# ---------------------------------------------------------------------------


class _TrackingSelector:
    """Strategy selector that records whether it was called."""

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

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


class _HalvingAllocator:
    """Budget allocator that halves the total budget for each candidate."""

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

    def allocate(
        self,
        candidates: list[tuple[Any, float]],
        total_budget: int,
        request: ContextRequest | None = None,
    ) -> list[tuple[Any, float, int]]:
        self.called = True
        half = total_budget // 2
        return [(s, c, half) for s, c in candidates]


class _ReversingExecutor:
    """Strategy executor that reverses the fragment order."""

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

    def execute(
        self,
        allocations: list[tuple[Any, float, int]],
        fragments: Sequence[ContextFragment],
        budget: ContextBudget,
    ) -> Sequence[ContextFragment]:
        self.called = True
        return list(reversed(fragments))


class _UkoDeduplicator:
    """Deduplicator that keeps only the first fragment per uko_node."""

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

    def deduplicate(
        self,
        fragments: Sequence[ContextFragment],
    ) -> Sequence[ContextFragment]:
        self.called = True
        seen: set[str] = set()
        result: list[ContextFragment] = []
        for f in fragments:
            if f.uko_node not in seen:
                seen.add(f.uko_node)
                result.append(f)
        return result


class _CustomPreamble:
    """Preamble generator that produces a known string."""

    def generate(
        self,
        fragments: Sequence[ContextFragment],
    ) -> str | None:
        return f"Custom preamble: {len(fragments)} fragments"


@given("a pipeline with a custom tracking strategy selector")
def step_pipeline_custom_selector(context: Context) -> None:
    selector = _TrackingSelector()
    context.custom_selector = selector
    context.pipeline = ACMSPipeline(strategy_selector=selector)


@given("a pipeline with a custom budget allocator that halves the budget")
def step_pipeline_custom_allocator(context: Context) -> None:
    allocator = _HalvingAllocator()
    context.custom_allocator = allocator
    context.pipeline = ACMSPipeline(budget_allocator=allocator)


@given("a pipeline with a custom strategy executor that reverses fragments")
def step_pipeline_custom_executor(context: Context) -> None:
    executor = _ReversingExecutor()
    context.custom_executor = executor
    context.pipeline = ACMSPipeline(strategy_executor=executor)


@given("a pipeline with a custom deduplicator that removes duplicates by uko_node")
def step_pipeline_custom_dedup(context: Context) -> None:
    dedup = _UkoDeduplicator()
    context.custom_deduplicator = dedup
    context.pipeline = ACMSPipeline(deduplicator=dedup)


@given("a pipeline with a custom preamble generator")
def step_pipeline_custom_preamble(context: Context) -> None:
    context.pipeline = ACMSPipeline(preamble_generator=_CustomPreamble())


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


@then("the custom budget allocator should have been called")
def step_allocator_called(context: Context) -> None:
    assert context.custom_allocator.called, "Custom allocator was not invoked"


@then("the custom strategy executor should have been called")
def step_executor_called(context: Context) -> None:
    assert context.custom_executor.called, "Custom executor was not invoked"


@then("the custom deduplicator should have been called")
def step_deduplicator_called(context: Context) -> None:
    assert context.custom_deduplicator.called, "Custom deduplicator was not invoked"


@then('the payload preamble should be "{expected}"')
def step_preamble_value(context: Context, expected: str) -> None:
    assert context.payload.preamble == expected, (
        f"Expected preamble {expected!r}, got {context.payload.preamble!r}"
    )


# ---------------------------------------------------------------------------
# T2 — SkeletonCompressor steps
# ---------------------------------------------------------------------------


@when("I compress fragments with the default skeleton compressor and budget {budget:d}")
def step_compress_default(context: Context, budget: int) -> None:
    frags = (
        _make_fragment(uko_node="test://a", content="hello", token_count=10),
        _make_fragment(uko_node="test://b", content="world", token_count=10),
    )
    context.original_fragments = frags
    compressor = DefaultSkeletonCompressor()
    context.compressed_fragments = compressor.compress(frags, budget)


@then("the compressed fragments should equal the original fragments")
def step_compressed_equal(context: Context) -> None:
    assert context.compressed_fragments == context.original_fragments


@when(
    "I compress fragments with a truncating skeleton compressor and budget {budget:d}"
)
def step_compress_truncating(context: Context, budget: int) -> None:
    frags = (
        _make_fragment(uko_node="test://a", content="a" * 100, token_count=100),
        _make_fragment(uko_node="test://b", content="b" * 100, token_count=100),
    )
    context.original_fragments = frags

    class _TruncatingCompressor:
        def compress(
            self,
            fragments: tuple[ContextFragment, ...],
            skeleton_budget: int,
        ) -> tuple[ContextFragment, ...]:
            result = []
            remaining = skeleton_budget
            for f in fragments:
                if remaining <= 0:
                    break
                # Re-create with truncated content.
                trimmed = f.content[:remaining]
                result.append(
                    _make_fragment(
                        uko_node=f.uko_node,
                        content=trimmed,
                        token_count=min(f.token_count, remaining),
                    )
                )
                remaining -= min(f.token_count, remaining)
            return tuple(result)

    compressor = _TruncatingCompressor()
    context.compressed_fragments = compressor.compress(frags, budget)


@then("the compressed fragments should have reduced content")
def step_compressed_reduced(context: Context) -> None:
    orig_total = sum(len(f.content) for f in context.original_fragments)
    comp_total = sum(len(f.content) for f in context.compressed_fragments)
    assert comp_total < orig_total, (
        f"Expected reduced content: original={orig_total}, compressed={comp_total}"
    )


# ---------------------------------------------------------------------------
# T3 — Multi-candidate allocation steps
# ---------------------------------------------------------------------------


@when(
    "I allocate budget {budget:d} across candidates with confidences {c1:g} and {c2:g}"
)
def step_allocate_multi(context: Context, budget: int, c1: float, c2: float) -> None:
    from cleveragents.application.services.acms_service import (
        RecencyStrategy,
        RelevanceStrategy,
    )

    s1 = RelevanceStrategy()
    s2 = RecencyStrategy()
    allocator = DefaultBudgetAllocator()
    context.allocations = allocator.allocate([(s1, c1), (s2, c2)], budget)
    context.alloc_budget = budget


@then("the first candidate should receive approximately {tokens:d} tokens")
def step_first_alloc(context: Context, tokens: int) -> None:
    actual = context.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")
def step_second_alloc(context: Context, tokens: int) -> None:
    actual = context.allocations[1][2]
    assert abs(actual - tokens) <= 1, (
        f"Second candidate got {actual}, expected ~{tokens}"
    )


@then("the total allocated should not exceed {budget:d}")
def step_total_alloc(context: Context, budget: int) -> None:
    total = sum(a[2] for a in context.allocations)
    assert total <= budget, f"Total allocated {total} exceeds budget {budget}"


@when("I select strategies with no strategy hint from {count:d} registered strategies")
def step_select_no_hint(context: Context, count: int) -> None:
    from cleveragents.application.services.acms_service import (
        RecencyStrategy,
        RelevanceStrategy,
        TieredStrategy,
    )

    strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
    selector = DefaultStrategySelector()
    context.selected = selector.select(strategies[:count], {})


@then("all {count:d} strategies should be returned with confidence scores")
def step_all_selected(context: Context, count: int) -> None:
    assert len(context.selected) == count, (
        f"Expected {count} selected strategies, got {len(context.selected)}"
    )
    for s, c in context.selected:
        assert c > 0.0, f"Strategy {s.name} has zero confidence"


# ---------------------------------------------------------------------------
# plan_id ULID validation steps (SEC-1, SPEC-1, TEST-GAP-1)
# ---------------------------------------------------------------------------


@when("I create a payload with an empty plan_id")
def step_create_payload_empty_plan_id(context: Context) -> None:
    context.acms_error = None
    context.test_payload = None
    try:
        context.test_payload = ContextPayload(plan_id="")
    except ValidationError as exc:
        context.acms_error = exc


@when('I create a payload with plan_id "{plan_id}"')
def step_create_payload_plan_id(context: Context, plan_id: str) -> None:
    context.acms_error = None
    context.test_payload = None
    try:
        context.test_payload = ContextPayload(plan_id=plan_id)
    except ValidationError as exc:
        context.acms_error = exc


@then('the created payload plan_id should be "{plan_id}"')
def step_created_payload_plan_id(context: Context, plan_id: str) -> None:
    assert context.test_payload is not None, "Payload was not created"
    assert context.test_payload.plan_id == plan_id, (
        f"Expected plan_id={plan_id!r}, got {context.test_payload.plan_id!r}"
    )


# ---------------------------------------------------------------------------
# Strategy introspection steps (capabilities, can_handle, explain)
# ---------------------------------------------------------------------------

_BUILTIN_STRATEGIES: dict[str, Any] = {
    "relevance": RelevanceStrategy(),
    "recency": RecencyStrategy(),
    "tiered": TieredStrategy(),
}


@when('I inspect the capabilities of "{strategy_name}"')
def step_inspect_capabilities(context: Context, strategy_name: str) -> None:
    context.inspected_capabilities = _BUILTIN_STRATEGIES[strategy_name].capabilities


@then("the strategy should declare {attr} as {expected}")
def step_capability_attr(context: Context, attr: str, expected: str) -> None:
    caps = context.inspected_capabilities
    actual = getattr(caps, attr)
    expected_val = expected == "true"
    assert actual == expected_val, f"Expected {attr}={expected_val}, got {actual}"


@when("I query can_handle on all three built-in strategies")
def step_query_can_handle(context: Context) -> None:
    request: dict[str, Any] = {"strategy": "relevance"}
    context.confidence_map = {
        name: s.can_handle(request) for name, s in _BUILTIN_STRATEGIES.items()
    }


@then('"{a}" should have the highest confidence')
def step_highest_confidence(context: Context, a: str) -> None:
    conf_a = context.confidence_map[a]
    for name, conf in context.confidence_map.items():
        if name != a:
            assert conf_a > conf, f"Expected {a} ({conf_a}) > {name} ({conf})"


@then('"{a}" should have higher confidence than "{b}"')
def step_higher_confidence(context: Context, a: str, b: str) -> None:
    assert context.confidence_map[a] > context.confidence_map[b], (
        f"Expected {a} ({context.confidence_map[a]}) > "
        f"{b} ({context.confidence_map[b]})"
    )


@when("I call explain on each built-in strategy")
def step_call_explain(context: Context) -> None:
    context.explanations = {
        name: s.explain() for name, s in _BUILTIN_STRATEGIES.items()
    }


@then('the "{strategy}" explanation should mention "{keyword}"')
def step_explain_keyword(context: Context, strategy: str, keyword: str) -> None:
    explanation = context.explanations[strategy]
    assert keyword.lower() in explanation.lower(), (
        f"Expected '{keyword}' in {strategy} explanation: {explanation!r}"
    )


# ---------------------------------------------------------------------------
# Fragment metadata overflow steps
# ---------------------------------------------------------------------------


@when("I create a fragment with {count:d} metadata entries")
def step_create_fragment_metadata(context: Context, count: int) -> None:
    context.acms_error = None
    context.test_fragment = None
    metadata = {f"key_{i}": f"val_{i}" for i in range(count)}
    try:
        context.test_fragment = ContextFragment(
            uko_node="project://test/meta.py",
            content="test",
            relevance_score=0.5,
            token_count=10,
            tier="hot",
            provenance=FragmentProvenance(resource_uri="project://test/meta.py"),
            metadata=metadata,
        )
    except ValidationError as exc:
        context.acms_error = exc


@then("the fragment should have {count:d} metadata entries")
def step_fragment_metadata_count(context: Context, count: int) -> None:
    assert context.test_fragment is not None, "Fragment was not created"
    assert len(context.test_fragment.metadata) == count


# ---------------------------------------------------------------------------
# Early plan_id rejection in assemble()
# ---------------------------------------------------------------------------


@when('I assemble with plan_id "{plan_id}"')
def step_assemble_with_plan_id(context: Context, plan_id: str) -> None:
    if not hasattr(context, "pipeline"):
        context.pipeline = ACMSPipeline()
    context.assemble_error = None
    try:
        context.payload = context.pipeline.assemble(
            plan_id=plan_id,
            fragments=list(context.fragments),
            budget=context.budget,
        )
    except (ValueError, ValidationError) as exc:
        context.assemble_error = exc


@then('an ACMS ValueError should be raised mentioning "{keyword}"')
def step_acms_value_error_mentioning(context: Context, keyword: str) -> None:
    assert context.assemble_error is not None, "No error was raised"
    assert isinstance(context.assemble_error, ValueError), (
        f"Expected ValueError, got {type(context.assemble_error).__name__}"
    )
    assert keyword.lower() in str(context.assemble_error).lower(), (
        f"Expected '{keyword}' in error: {context.assemble_error}"
    )


# ---------------------------------------------------------------------------
# TEST-2 — Provenance map immutability step
# ---------------------------------------------------------------------------


@when("I create a payload with a mutable provenance map and mutate the source")
def step_create_payload_mutable_provenance(context: Context) -> None:
    inner = {"strategy": "relevance", "extra": "value"}
    source_map = {"frag-1": inner}
    context.test_payload = ContextPayload(
        plan_id="01JQTESTPN00000000000000AB",
        provenance_map=source_map,
    )
    # Mutate the source dict after construction.
    inner["strategy"] = "TAMPERED"
    source_map["frag-2"] = {"injected": "true"}


@then("the payload provenance map should be unchanged")
def step_provenance_map_unchanged(context: Context) -> None:
    pmap = context.test_payload.provenance_map
    assert "frag-1" in pmap, "frag-1 missing from provenance_map"
    assert "frag-2" not in pmap, "frag-2 was injected into provenance_map"
    assert pmap["frag-1"]["strategy"] == "relevance", (
        f"Expected 'relevance', got {pmap['frag-1']['strategy']!r} — "
        "external mutation leaked into frozen model"
    )


# ---------------------------------------------------------------------------
# TEST-3 — Allocator rounding edge case steps
# ---------------------------------------------------------------------------


@when("I allocate budget {budget:d} across {count:d} candidates with equal confidence")
def step_allocate_equal(context: Context, budget: int, count: int) -> None:
    from cleveragents.application.services.acms_service import (
        RecencyStrategy,
        RelevanceStrategy,
        TieredStrategy,
    )

    strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
    allocator = DefaultBudgetAllocator()
    candidates = [(s, 0.5) for s in strategies[:count]]
    context.allocations = allocator.allocate(candidates, budget)
    context.alloc_budget = budget


@when("I allocate budget {budget:d} across {count:d} candidates with zero confidence")
def step_allocate_zero_conf(context: Context, budget: int, count: int) -> None:
    from cleveragents.application.services.acms_service import (
        RecencyStrategy,
        RelevanceStrategy,
        TieredStrategy,
    )

    strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
    allocator = DefaultBudgetAllocator()
    candidates = [(s, 0.0) for s in strategies[:count]]
    context.allocations = allocator.allocate(candidates, budget)
    context.alloc_budget = budget


@when(
    "I allocate budget {budget:d} across {count:d} candidates with confidences {conf_str}"
)
def step_allocate_unequal(
    context: Context, budget: int, count: int, conf_str: str
) -> None:
    from cleveragents.application.services.acms_service import (
        RecencyStrategy,
        RelevanceStrategy,
        TieredStrategy,
    )

    confidences = [float(c) for c in conf_str.split()]
    strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
    allocator = DefaultBudgetAllocator()
    candidates = [(s, c) for s, c in zip(strategies[:count], confidences, strict=True)]
    context.allocations = allocator.allocate(candidates, budget)
    context.alloc_budget = budget


@then("the total allocated should equal exactly {budget:d}")
def step_total_exact(context: Context, budget: int) -> None:
    total = sum(a[2] for a in context.allocations)
    assert total == budget, (
        f"Expected total allocation to be exactly {budget}, got {total}"
    )


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


# ---------------------------------------------------------------------------
# TEST-4 — Strategy selector fallback step
# ---------------------------------------------------------------------------


class _EmptySelector:
    """Strategy selector that always returns an empty candidate list."""

    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 []


@given("a pipeline with a custom strategy selector that returns no candidates")
def step_pipeline_empty_selector(context: Context) -> None:
    selector = _EmptySelector()
    context.empty_selector = selector
    context.pipeline = ACMSPipeline(strategy_selector=selector)


@then("the empty strategy selector should have been called")
def step_empty_selector_called(context: Context) -> None:
    assert context.empty_selector.called, (
        "Expected _EmptySelector.select() to be called, but it was not"
    )


# ---------------------------------------------------------------------------
# Issue #3563 — SkeletonCompressor invocation in assemble() steps
# ---------------------------------------------------------------------------


class _SpySkeletonCompressor:
    """Spy skeleton compressor that records calls and returns fragments unchanged."""

    def __init__(self) -> None:
        self.called = False
        self.received_fragments: tuple[ContextFragment, ...] = ()
        self.received_budget: int = 0
        self.output: tuple[ContextFragment, ...] = ()

    def compress(
        self,
        fragments: tuple[ContextFragment, ...],
        skeleton_budget: int,
    ) -> tuple[ContextFragment, ...]:
        self.called = True
        self.received_fragments = fragments
        self.received_budget = skeleton_budget
        # Return the fragments unchanged (pass-through spy)
        self.output = fragments
        return self.output


@given("a spy skeleton compressor is registered")
def step_spy_compressor_registered(context: Context) -> None:
    spy = _SpySkeletonCompressor()
    context.spy_compressor = spy
    context.pipeline = ACMSPipeline(skeleton_compressor=spy)


@given("parent fragments for skeleton compression:")
def step_parent_fragments(context: Context) -> None:
    frags = []
    for row in context.table:
        frags.append(
            _make_fragment(
                uko_node=row["uko_node"],
                content=row["content"],
                relevance_score=float(row["score"]),
                token_count=int(row["tokens"]),
            )
        )
    context.parent_fragments = tuple(frags)


@when("I assemble with parent fragments and skeleton_ratio {ratio:g}")
def step_assemble_with_parent_fragments(context: Context, ratio: float) -> None:
    if not hasattr(context, "pipeline"):
        context.pipeline = ACMSPipeline()
    context.assemble_error = None
    try:
        context.payload = context.pipeline.assemble(
            plan_id="01JQTESTPN00000000000000AB",
            fragments=list(context.fragments),
            budget=context.budget,
            parent_fragments=context.parent_fragments,
            skeleton_ratio=ratio,
        )
    except (ValueError, ValidationError) as exc:
        context.assemble_error = exc


@then("the spy skeleton compressor should have been called")
def step_spy_compressor_called(context: Context) -> None:
    assert context.spy_compressor.called, (
        "Expected _SpySkeletonCompressor.compress() to be called, but it was not. "
        "This confirms the bug: SkeletonCompressor is stored but never invoked in assemble()."
    )


@then("the spy compressor should have received the parent fragments")
def step_spy_compressor_received_fragments(context: Context) -> None:
    assert context.spy_compressor.received_fragments == context.parent_fragments, (
        f"Expected compressor to receive parent_fragments, "
        f"got {context.spy_compressor.received_fragments!r}"
    )


@then("the spy compressor skeleton_budget should be {expected_budget:d}")
def step_spy_compressor_budget(context: Context, expected_budget: int) -> None:
    actual = context.spy_compressor.received_budget
    assert actual == expected_budget, (
        f"Expected skeleton_budget={expected_budget}, got {actual}. "
        f"skeleton_budget should be int(total_budget * skeleton_ratio)."
    )


@then("the payload skeleton_fragments should be non-empty")
def step_payload_skeleton_non_empty(context: Context) -> None:
    assert context.payload.skeleton_fragments, (
        "Expected payload.skeleton_fragments to be non-empty, but it was empty. "
        "The assembled context must include compressed skeleton fragments."
    )


@then("the payload skeleton_fragments should equal the spy compressor output")
def step_payload_skeleton_equals_spy_output(context: Context) -> None:
    assert context.payload.skeleton_fragments == context.spy_compressor.output, (
        f"Expected payload.skeleton_fragments == spy output, "
        f"got {context.payload.skeleton_fragments!r}"
    )


@then("the payload skeleton_fragments should be empty")
def step_payload_skeleton_empty(context: Context) -> None:
    assert context.payload.skeleton_fragments == (), (
        f"Expected payload.skeleton_fragments to be empty (), "
        f"got {context.payload.skeleton_fragments!r}"
    )