"""Tests for model routing learning."""

from __future__ import annotations

import pytest

from logos_adk.agent import Agent
from logos_adk.learning import ModelRoutingLearningEngine, RunRecord, SQLiteLearningStore
from logos_adk.providers.mock import MockProvider
from logos_adk.runtime import _default_runtime_holder


@pytest.fixture(autouse=True)
def _reset_runtime():
    _default_runtime_holder.reset()
    yield
    _default_runtime_holder.reset()


@pytest.mark.asyncio
async def test_per_task_model_routing_changes_selected_provider(tmp_path):
    engine = ModelRoutingLearningEngine(path=str(tmp_path / "routing.db"))
    cheap = MockProvider(responses=["cheap response"])
    cheap.model_name = "gpt-4o-mini"
    strong = MockProvider(responses=["strong response"])
    strong.model_name = "claude-sonnet-4-20250514"

    for _ in range(4):
        engine.record_route(
            model_key="MockProvider:gpt-4o-mini",
            success=True,
            latency_ms=25,
            cost_usd=0.0002,
            confidence_tier="low",
            workspace_id="acme",
            task_type="support",
            quality_score=0.3,
        )
        engine.record_route(
            model_key="MockProvider:claude-sonnet-4-20250514",
            success=True,
            latency_ms=35,
            cost_usd=0.003,
            confidence_tier="high",
            workspace_id="acme",
            task_type="analysis",
            quality_score=0.95,
        )

    agent = (
        Agent("router")
        .model(cheap)
        .fallback(strong)
        .model_routing_learning(engine)
    )

    support = await agent.run("Need a short support reply", _workspace_id="acme", _task_type="support")
    analysis = await agent.run(
        "Analyze the architectural tradeoffs of this migration plan in detail",
        _workspace_id="acme",
        _task_type="analysis",
    )

    assert support.content == "cheap response"
    assert analysis.content == "strong response"


@pytest.mark.asyncio
async def test_budget_aware_routing_prefers_cheapest_model_near_limit(tmp_path):
    engine = ModelRoutingLearningEngine(path=str(tmp_path / "budget.db"))
    cheap = MockProvider(responses=["cheap route"])
    cheap.model_name = "gpt-4o-mini"
    strong = MockProvider(responses=["strong route"])
    strong.model_name = "claude-sonnet-4-20250514"

    for _ in range(4):
        engine.record_route(
            model_key="MockProvider:claude-sonnet-4-20250514",
            success=True,
            latency_ms=40,
            cost_usd=0.006,
            confidence_tier="high",
            workspace_id="acme",
            task_type="analysis",
            quality_score=0.98,
        )
    engine.set_budget(
        workspace_id="acme",
        task_type="analysis",
        soft_limit_usd=0.02,
        threshold_ratio=0.8,
    )

    agent = (
        Agent("router")
        .model(cheap)
        .fallback(strong)
        .model_routing_learning(engine)
    )

    result = await agent.run(
        "Analyze this migration design and call out the hidden risks",
        _workspace_id="acme",
        _task_type="analysis",
    )
    assert result.content == "cheap route"

    status = engine.budget_status(workspace_id="acme", task_type="analysis")
    assert status["near_limit"] is True


def test_confidence_estimation_and_model_leaderboard(tmp_path):
    engine = ModelRoutingLearningEngine(path=str(tmp_path / "confidence.db"))
    for _ in range(3):
        engine.record_route(
            model_key="MockProvider:gpt-4o-mini",
            success=True,
            latency_ms=18,
            cost_usd=0.0002,
            confidence_tier="low",
            workspace_id="acme",
            task_type="support",
            quality_score=0.4,
        )
        engine.record_route(
            model_key="MockProvider:claude-sonnet-4-20250514",
            success=True,
            latency_ms=42,
            cost_usd=0.003,
            confidence_tier="high",
            workspace_id="acme",
            task_type="analysis",
            quality_score=0.9,
        )

    assert engine.estimate_confidence("hello", task_type="support") == "low"
    assert engine.estimate_confidence("Please analyze the architecture in detail", task_type="analysis") == "high"

    support_scores = engine.list_model_scores(workspace_id="acme", task_type="support", confidence_tier="low")
    analysis_scores = engine.list_model_scores(workspace_id="acme", task_type="analysis", confidence_tier="high")
    assert support_scores[0].model_key == "MockProvider:gpt-4o-mini"
    assert analysis_scores[0].model_key == "MockProvider:claude-sonnet-4-20250514"


def test_sync_trace_updates_quality_from_feedback_and_outcomes(tmp_path):
    engine = ModelRoutingLearningEngine(path=str(tmp_path / "sync.db"))
    store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    store.create_run(
        RunRecord(
            id="run-1",
            trace_id="trace-1",
            request_id="req-1",
            agent_name="router",
            input_text="question",
            output_text="answer",
            workspace_id="acme",
            model="claude-sonnet-4-20250514",
            config_fingerprint="cfg",
        )
    )
    engine.record_route(
        model_key="MockProvider:claude-sonnet-4-20250514",
        success=True,
        latency_ms=30,
        cost_usd=0.002,
        confidence_tier="high",
        workspace_id="acme",
        task_type="analysis",
        trace_id="trace-1",
    )

    from logos_adk.learning import FeedbackAPI

    api = FeedbackAPI(store)
    api.record_feedback(trace_id="trace-1", kind="thumbs_up")
    api.record_outcome(trace_id="trace-1", kind="conversion", value=True)
    engine.sync_trace("trace-1", store)

    scores = engine.list_model_scores(workspace_id="acme", task_type="analysis", confidence_tier="high")
    assert scores[0].average_quality_score > 0.5


def test_automatic_canarying_routes_partial_traffic(tmp_path):
    engine = ModelRoutingLearningEngine(path=str(tmp_path / "canary.db"))
    engine.register_canary(
        model_key="MockProvider:new-model",
        rollout_percentage=50,
        workspace_id="acme",
        task_type="support",
    )

    used_canary = False
    skipped_canary = False
    for index in range(64):
        decision = engine.route_candidates(
            ["MockProvider:baseline", "MockProvider:new-model"],
            prompt="Need a short support reply",
            assignment_key=f"session-{index}",
            workspace_id="acme",
            task_type="support",
        )
        if decision.used_canary:
            used_canary = True
        else:
            skipped_canary = True
        if used_canary and skipped_canary:
            break

    assert used_canary is True
    assert skipped_canary is True
