"""Tests for Agent.stream() and streaming chunks."""
import pytest

from logos_adk.agent import Agent
from logos_adk.providers import provider_registry
from logos_adk.providers.mock import MockProvider
from logos_adk.runtime import _default_runtime_holder
from logos_adk.tools import tool
from logos_adk.types import Response, StreamDone, TextChunk, ToolCall, ToolEnd, ToolStart, Usage


@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()


@pytest.mark.asyncio
async def test_agent_stream_yields_text_chunks_and_done():
    provider = MockProvider(stream_responses=[["Hel", "lo"]])
    agent = Agent("bot").model(provider)

    chunks = [chunk async for chunk in agent.stream("hi")]

    assert [chunk.text for chunk in chunks if isinstance(chunk, TextChunk)] == ["Hel", "lo"]
    assert isinstance(chunks[-1], StreamDone)
    assert chunks[-1].response.content == "Hello"


@pytest.mark.asyncio
async def test_agent_stream_handles_tool_loop():
    @tool
    async def add(a: int, b: int) -> int:
        return a + b

    tc = ToolCall(id="tc1", name="add", arguments={"a": 2, "b": 3})
    provider = MockProvider(
        stream_responses=[
            Response(content="", tool_calls=[tc], usage=Usage(1, 1), raw={}),
            ["5", " total"],
        ]
    )
    agent = Agent("bot").model(provider).tools([add])

    chunks = [chunk async for chunk in agent.stream("compute")]

    assert any(isinstance(chunk, ToolStart) and chunk.name == "add" for chunk in chunks)
    assert any(
        isinstance(chunk, ToolEnd)
        and chunk.result.content == 5
        and chunk.result.error is None
        for chunk in chunks
    )
    assert [chunk.text for chunk in chunks if isinstance(chunk, TextChunk)] == ["5", " total"]
    assert isinstance(chunks[-1], StreamDone)
    assert chunks[-1].response.content == "5 total"


@pytest.mark.asyncio
async def test_agent_stream_updates_session_state():
    provider = MockProvider(stream_responses=[["Hi"]])
    agent = Agent("bot").model(provider)

    chunks = [chunk async for chunk in agent.stream("hello", session_id="s1")]
    assert isinstance(chunks[-1], StreamDone)

    state = agent._get_state("s1")
    assert [message.content for message in state.messages] == ["hello", "Hi"]


@pytest.mark.asyncio
async def test_agent_stream_events_include_user_facing_progress():
    provider = MockProvider(stream_responses=[["Hi"]])
    agent = Agent("bot").model(provider)

    events = [event async for event in agent.stream_events("hello", session_id="s1")]

    assert events[0]["event"] == "graph_started"
    assert events[0]["user_facing"]["kind"] == "run"
    assert events[1]["event"] == "node_started"
    assert events[1]["user_facing"]["kind"] == "step"
    assert events[-1]["event"] == "graph_completed"
    assert events[-1]["user_facing"]["kind"] == "result"


@pytest.mark.asyncio
async def test_agent_run_invokes_text_and_done_callbacks():
    provider = MockProvider(stream_responses=[["Hel", "lo"]])
    agent = Agent("bot").model(provider)
    seen_text: list[str] = []
    seen_done: list[str] = []

    result = await agent.run(
        "hi",
        on_text=lambda chunk: seen_text.append(chunk.text),
        on_done=lambda response: seen_done.append(response.content),
    )

    assert result.content == "Hello"
    assert seen_text == ["Hel", "lo"]
    assert seen_done == ["Hello"]


@pytest.mark.asyncio
async def test_agent_run_invokes_tool_callbacks():
    @tool
    async def add(a: int, b: int) -> int:
        return a + b

    tc = ToolCall(id="tc1", name="add", arguments={"a": 2, "b": 3})
    provider = MockProvider(
        stream_responses=[
            Response(content="", tool_calls=[tc], usage=Usage(1, 1), raw={}),
            ["5"],
        ]
    )
    agent = Agent("bot").model(provider).tools([add])
    seen_events: list[tuple[str, str, int | None]] = []

    result = await agent.run(
        "compute",
        on_tool_start=lambda chunk: seen_events.append(("start", chunk.name, None)),
        on_tool_end=lambda chunk: seen_events.append(
            ("end", chunk.result.call_id, chunk.result.content)
        ),
    )

    assert result.content == "5"
    assert seen_events == [("start", "add", None), ("end", "tc1", 5)]


@pytest.mark.asyncio
async def test_agent_run_supports_async_callbacks():
    provider = MockProvider(stream_responses=[["Hi"]])
    agent = Agent("bot").model(provider)
    seen: list[str] = []

    async def on_text(chunk: TextChunk) -> None:
        seen.append(chunk.text)

    result = await agent.run("hello", on_text=on_text)

    assert result.content == "Hi"
    assert seen == ["Hi"]


@pytest.mark.asyncio
async def test_agent_stream_buffers_until_self_qa_passes():
    provider = MockProvider(
        responses=[
            "Weak answer.",
            '{"assertions":[{"assertion":"The answer must name the exact action."},{"assertion":"The answer must be actionable."}]}',
            '{"passed":false,"summary":"Too vague.","checks":[{"assertion":"The answer must name the exact action.","passed":false,"failure_reason":"No exact action."},{"assertion":"The answer must be actionable.","passed":false,"failure_reason":"No concrete step."}]}',
            "Restart the API process and re-run the migration.",
            '{"assertions":[{"assertion":"The answer must name the exact action."},{"assertion":"The answer must be actionable."}]}',
            '{"passed":true,"summary":"Looks good.","checks":[{"assertion":"The answer must name the exact action.","passed":true,"evidence":"It names restart and migration."},{"assertion":"The answer must be actionable.","passed":true,"evidence":"It gives explicit steps."}]}',
        ]
    )
    agent = Agent("bot").model(provider).self_qa(enabled=True, max_retries=1)

    chunks = [chunk async for chunk in agent.stream("Fix the broken deployment")]

    assert [chunk.text for chunk in chunks if isinstance(chunk, TextChunk)] == [
        "Restart the API process and re-run the migration."
    ]
    assert isinstance(chunks[-1], StreamDone)
    assert chunks[-1].response.status == "completed"
    assert chunks[-1].response.raw["self_qa"]["passed"] is True


@pytest.mark.asyncio
async def test_agent_stream_buffers_in_strict_budget_mode_and_fails_closed():
    provider = MockProvider(
        responses=[
            Response(content="Expensive answer", tool_calls=[], usage=Usage(400, 200), raw={}),
        ]
    )
    provider.model_name = "claude-sonnet-4-20250514"
    agent = Agent("bot").model(provider).budget(task_budget_usd=0.001, fail_on_budget_exhausted=True)

    chunks = [chunk async for chunk in agent.stream("Analyze the migration plan")]

    text_chunks = [chunk.text for chunk in chunks if isinstance(chunk, TextChunk)]
    assert len(text_chunks) == 1
    assert text_chunks[0].startswith("Budget exhausted during llm_generate:")
    assert isinstance(chunks[-1], StreamDone)
    assert chunks[-1].response.status == "failed"
    assert chunks[-1].response.raw["budget_failure"]["error_type"] == "BudgetExhaustedError"

