from behave import given, then, when
from context_engine_registry import RegistryBuilder

from biblicus.context import ContextPack, ContextPackBlock
from biblicus.context_engine import ContextAssembler


@given("an explicit Context with prioritized packs and tight budget")
def step_explicit_context_priority(context):
    builder = RegistryBuilder()
    builder.register_corpus(
        "wikitext",
        {"source": "wikitext2", "split": "train", "maximum_cache_total_items": 2000},
    )
    builder.register_retriever(
        "search_primary",
        {
            "corpus": "wikitext",
            "query": "Valkyria Chronicles III",
            "limit": 5,
            "maximum_total_characters": 800,
        },
    )
    builder.register_retriever(
        "search_secondary",
        {
            "corpus": "wikitext",
            "query": "Valkyria Chronicles III",
            "limit": 5,
            "maximum_total_characters": 800,
        },
    )
    builder.register_retriever(
        "search_tertiary",
        {
            "corpus": "wikitext",
            "query": "Valkyria Chronicles III",
            "limit": 5,
            "maximum_total_characters": 800,
        },
    )
    builder.register_context(
        "support_context",
        {
            "policy": {
                "input_budget": {"max_tokens": 10},
                "pack_budget": {"default_ratio": 1.0},
                "overflow": "compact",
                "max_iterations": 2,
            },
            "messages": [
                {"type": "context", "name": "search_primary", "priority": 2},
                {"type": "context", "name": "search_secondary", "priority": 1},
                {"type": "context", "name": "search_tertiary", "priority": 0},
                {"type": "user", "content": "Question"},
            ],
        },
    )
    context.registry = builder.registry
    context.context_name = "support_context"
    context.calls = []


@when("I assemble that explicit Context with regeneration for priority")
def step_assemble_explicit_context(context):
    def fake_retrieve(request):
        context.calls.append(request.maximum_total_characters)
        text = "Valkyria Chronicles III long excerpt" * 5
        return ContextPack(
            text=text,
            evidence_count=1,
            blocks=[ContextPackBlock(evidence_item_id="fake-1", text=text, metadata=None)],
        )

    assembler = ContextAssembler(
        context.registry.contexts,
        retriever_registry=context.registry.retrievers,
        corpus_registry=context.registry.corpora,
    )
    assembler.assemble(
        context_name=context.context_name,
        base_system_prompt="",
        history_messages=[],
        user_message="",
        template_context={"input": {}, "context": {}},
        retriever_override=fake_retrieve,
    )


@then("higher-priority explicit packs should receive greater or equal budgets")
def step_priority_budgets(context):
    assert len(context.calls) >= 3
    assert context.calls[0] >= context.calls[1]
    assert context.calls[1] >= context.calls[2]
