"""Tests for tool and workflow learning."""

from __future__ import annotations

import pytest

from logos_adk.agent import Agent
from logos_adk.context import RunContext, use_run_context
from logos_adk.errors import ProviderError
from logos_adk.learning import ToolWorkflowLearningEngine
from logos_adk.providers import Provider, provider_registry
from logos_adk.providers.mock import MockProvider
from logos_adk.runtime import _default_runtime_holder
from logos_adk.state_store import InMemoryStateStore
from logos_adk.tools import tool
from logos_adk.types import Message, Response, Usage
from logos_adk.workflow.pipeline import Pipeline


@pytest.fixture(autouse=True)
def _reset_runtime_and_registry():
    _default_runtime_holder.reset()
    provider_registry._entries.clear()
    yield
    _default_runtime_holder.reset()
    provider_registry._entries.clear()


@tool
async def add(a: int, b: int) -> int:
    """Add two numbers."""
    return a + b


@tool
async def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b


class _FailingProvider(Provider):
    def __init__(self, model_name: str) -> None:
        self.model_name = model_name

    async def generate(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        output_schema: object = None,
        **kwargs: object,
    ) -> Response:
        raise ProviderError("boom")

    async def stream(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        **kwargs: object,
    ):
        raise ProviderError("boom")
        yield  # pragma: no cover


@pytest.mark.asyncio
async def test_tool_scoring_and_runtime_ordering(tmp_path):
    engine = ToolWorkflowLearningEngine(path=str(tmp_path / "tool_workflow.db"))
    engine.record_tool_call(tool_name="multiply", success=True, duration_ms=10, task_type="math")
    engine.record_tool_call(tool_name="multiply", success=True, duration_ms=12, task_type="math")
    engine.record_tool_call(tool_name="add", success=False, duration_ms=30, task_type="math")

    provider = MockProvider(responses=["done"])
    agent = Agent("bot").model(provider).tools([add, multiply]).tool_workflow_learning(engine)

    await agent.run("pick the best tool", _task_type="math")

    offered_tools = provider.call_history[0]["tools"]
    assert [tool_schema["name"] for tool_schema in offered_tools] == ["multiply", "add"]


@pytest.mark.asyncio
async def test_tool_calls_record_success_scores(tmp_path):
    from logos_adk.types import ToolCall

    engine = ToolWorkflowLearningEngine(path=str(tmp_path / "tool_calls.db"))
    provider = MockProvider(
        responses=[
            Response(
                content="",
                tool_calls=[ToolCall(id="tc1", name="multiply", arguments={"a": 3, "b": 4})],
                usage=Usage(0, 0),
                raw={},
            ),
            "twelve",
        ]
    )
    agent = Agent("bot").model(provider).tools([add, multiply]).tool_workflow_learning(engine)

    result = await agent.run("compute", _task_type="math", _workspace_id="acme")
    assert result.content == "twelve"

    scores = engine.list_tool_scores(task_type="math", workspace_id="acme")
    assert len(scores) == 1
    assert scores[0].tool_name == "multiply"
    assert scores[0].success_count == 1
    assert scores[0].failure_count == 0


@pytest.mark.asyncio
async def test_fallback_learning_reorders_fallback_candidates(tmp_path):
    engine = ToolWorkflowLearningEngine(path=str(tmp_path / "fallback.db"))
    engine.record_fallback(
        provider_key="MockProvider:slow",
        success=False,
        duration_ms=80,
        task_type="support",
    )
    engine.record_fallback(
        provider_key="MockProvider:fast",
        success=True,
        duration_ms=15,
        task_type="support",
    )

    primary = MockProvider(responses=["ok"])
    primary.model_name = "primary"
    fallback_slow = MockProvider(responses=["slow"])
    fallback_slow.model_name = "slow"
    fallback_fast = MockProvider(responses=["fast"])
    fallback_fast.model_name = "fast"
    agent = (
        Agent("bot")
        .model(primary)
        .fallback(fallback_slow, fallback_fast)
        .tool_workflow_learning(engine)
    )

    with use_run_context(RunContext(task_type="support")):
        ordered = agent._provider_candidates()

    assert ordered[0] is primary
    assert ordered[1] is fallback_fast
    assert ordered[2] is fallback_slow


@pytest.mark.asyncio
async def test_fallback_events_record_from_agent_runtime(tmp_path):
    engine = ToolWorkflowLearningEngine(path=str(tmp_path / "fallback_runtime.db"))
    primary = _FailingProvider("primary")
    fallback = MockProvider(responses=["recovered"])
    fallback.model_name = "backup"
    agent = Agent("bot").model(primary).fallback(fallback).tool_workflow_learning(engine)

    result = await agent.run("hi", _task_type="ops", _workspace_id="acme")
    assert result.content == "recovered"

    scores = engine.list_fallback_scores(task_type="ops", workspace_id="acme")
    assert len(scores) == 1
    assert scores[0].provider_key == "MockProvider:backup"
    assert scores[0].success_count == 1


@pytest.mark.asyncio
async def test_workflow_branch_loop_and_approval_learning(tmp_path):
    engine = ToolWorkflowLearningEngine(path=str(tmp_path / "workflow.db"))
    store = InMemoryStateStore()

    branch_pipeline = Pipeline("branching", learning_engine=engine)

    async def setup(ctx):
        ctx["flag"] = True
        return "setup"

    async def yes(ctx):
        return "yes"

    branch_pipeline.function(setup, name="setup").conditional(
        lambda ctx: ctx["flag"],
        then=yes,
        name="branch",
    )
    branch_response = await branch_pipeline.run("", _workspace_id="acme")
    assert branch_response.content == "yes"

    branch_scores = engine.list_workflow_branch_scores(
        pipeline_name="branching",
        step_name="branch",
        workspace_id="acme",
    )
    assert len(branch_scores) == 1
    assert branch_scores[0].branch_name == "yes"
    assert branch_scores[0].success_count == 1

    for _ in range(4):
        engine.record_loop_execution(
            pipeline_name="looping",
            step_name="counter",
            iterations=2,
            max_iterations=10,
            termination_reason="condition_met",
            success=True,
            workspace_id="acme",
        )

    loop_pipeline = Pipeline("looping", learning_engine=engine)

    async def again(ctx):
        ctx["count"] = ctx.get("count", 0) + 1
        return ctx["count"]

    loop_pipeline.loop(again, until=lambda ctx: False, max_iterations=10, name="counter")
    loop_response = await loop_pipeline.run("", _workspace_id="acme")
    counter_meta = loop_response.raw["pipeline_result"]["outputs"]["counter"]["metadata"]
    assert counter_meta["max_iterations"] == 3
    assert counter_meta["iterations"] == 3

    approval_pipeline = Pipeline("approvaling", state_store=store, learning_engine=engine)
    approval_pipeline.function(setup, name="before").approval(name="review", message="Review this")
    paused = await approval_pipeline.run("", _workspace_id="acme")
    pipeline_id = paused.raw["pipeline_result"]["pipeline_id"]
    result = await approval_pipeline.resume(pipeline_id, approval=False)
    assert result.status == "failed"

    approval_scores = engine.list_approval_scores(
        pipeline_name="approvaling",
        step_name="review",
        workspace_id="acme",
    )
    assert len(approval_scores) == 1
    assert approval_scores[0].request_count == 1
    assert approval_scores[0].rejection_count == 1
    assert approval_scores[0].correction_rate == pytest.approx(1.0)
