from ai_orchestrator.application.services.router import (
    ConfigurableModelRouter,
    ProviderResponse,
)
from ai_orchestrator.config import (
    ExecutionConfig,
    McpServerConfig,
    ModelsConfig,
    ModelSlotConfig,
    PolicyConfig,
    ProjectConfig,
    ProjectMetadata,
)
from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider
from ai_orchestrator.infrastructure.storage.memory import InMemoryInvocationStore


def test_model_router_retries_weak_and_falls_back_to_strong() -> None:
    config = ProjectConfig(
        project=ProjectMetadata(id="default", name="Default"),
        models=ModelsConfig(
            weak=ModelSlotConfig(provider="local", model="weak"),
            strong=ModelSlotConfig(provider="external", model="strong"),
            vision=ModelSlotConfig(provider="disabled"),
            embedding=ModelSlotConfig(provider="local", model="embed"),
        ),
        mcp_servers={"test": McpServerConfig(base_url="http://example.test")},
        execution=ExecutionConfig(allow_external_models=True),
        policy=PolicyConfig(),
    )
    invocation_store = InMemoryInvocationStore()
    router = ConfigurableModelRouter(
        project_configs={"default": config},
        providers={
            "local": StaticMockModelProvider(
                provider_name="local",
                responses=[
                    ProviderResponse(
                        status="success",
                        message={"role": "assistant", "content": ""},
                        tool_calls=[],
                        usage={},
                    ),
                    ProviderResponse(
                        status="success",
                        message={"role": "assistant", "content": ""},
                        tool_calls=[],
                        usage={},
                    ),
                ],
            ),
            "external": StaticMockModelProvider(
                provider_name="external",
                responses=[
                    ProviderResponse(
                        status="success",
                        message={"role": "assistant", "content": "strong answer"},
                        tool_calls=[],
                        usage={"input_tokens": 1, "output_tokens": 2},
                    )
                ],
            ),
        },
        invocation_store=invocation_store,
    )

    result = router.run(
        project_id="default",
        slot="weak",
        messages=[{"role": "user", "content": "Hello"}],
        task_context={"task_id": "task_1", "node_id": "node_1"},
    )

    assert result.status == "success"
    assert result.model == "strong"
    assert result.message["content"] == "strong answer"
    assert len(invocation_store.list_model_invocations("task_1")) == 3