"""Integration tests using the A2A client SDK against the ASGI app."""

from __future__ import annotations

from collections.abc import AsyncIterator
from typing import Any

import httpx
import pytest
from a2a.client import A2ACardResolver, ClientConfig, create_client
from a2a.helpers import get_stream_response_text
from a2a.helpers.proto_helpers import new_text_message
from a2a.types.a2a_pb2 import Role, SendMessageRequest
from fastapi import FastAPI
from langchain_core.messages import AIMessage
from langgraph.graph import END, START, MessagesState, StateGraph

from a2a_langgraph.fastapi import add_langgraph_fastapi_endpoint

_MOUNT = "/agent"
_BASE_URL = "http://test"
_AGENT_BASE = f"{_BASE_URL}{_MOUNT}"


def _echo_graph() -> Any:
    def chatbot(state: MessagesState) -> dict[str, Any]:
        last = state["messages"][-1]
        return {"messages": [AIMessage(content=f"echo:{last.content}")]}

    builder = StateGraph(MessagesState)
    builder.add_node("chatbot", chatbot)
    builder.add_edge(START, "chatbot")
    builder.add_edge("chatbot", END)
    return builder.compile()


@pytest.fixture
def a2a_app() -> FastAPI:
    app = FastAPI()
    add_langgraph_fastapi_endpoint(
        app,
        graph=_echo_graph(),
        mount_path=_MOUNT,
        base_url=_BASE_URL,
        agent_name="Test Agent",
        agent_description="Integration test",
    )
    return app


@pytest.fixture
async def asgi_client(a2a_app: FastAPI) -> AsyncIterator[httpx.AsyncClient]:
    transport = httpx.ASGITransport(app=a2a_app)
    async with httpx.AsyncClient(transport=transport, base_url=_BASE_URL) as client:
        yield client


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _collect_text(chunks: list[Any]) -> str:
    return " ".join(
        t for chunk in chunks if (t := get_stream_response_text(chunk))
    )


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_card_resolver(asgi_client: httpx.AsyncClient) -> None:
    """A2ACardResolver should retrieve and parse the public agent card."""
    resolver = A2ACardResolver(httpx_client=asgi_client, base_url=_AGENT_BASE)
    card = await resolver.get_agent_card()

    assert card.name == "Test Agent"
    assert card.url == _AGENT_BASE


@pytest.mark.asyncio
async def test_non_streaming_send(asgi_client: httpx.AsyncClient) -> None:
    """Non-streaming client should return the full echo response in one pass."""
    resolver = A2ACardResolver(httpx_client=asgi_client, base_url=_AGENT_BASE)
    card = await resolver.get_agent_card()

    config = ClientConfig(streaming=False, httpx_client=asgi_client)
    client = await create_client(agent=card, client_config=config)
    try:
        request = SendMessageRequest(
            message=new_text_message("hello", role=Role.ROLE_USER)
        )
        chunks = [chunk async for chunk in client.send_message(request)]
        assert chunks, "Expected at least one response chunk"
        assert "echo:hello" in _collect_text(chunks)
    finally:
        await client.close()


@pytest.mark.asyncio
async def test_streaming_send(asgi_client: httpx.AsyncClient) -> None:
    """Streaming client should yield chunks whose text assembles the echo reply."""
    resolver = A2ACardResolver(httpx_client=asgi_client, base_url=_AGENT_BASE)
    card = await resolver.get_agent_card()

    config = ClientConfig(streaming=True, httpx_client=asgi_client)
    client = await create_client(agent=card, client_config=config)
    try:
        request = SendMessageRequest(
            message=new_text_message("stream me", role=Role.ROLE_USER)
        )
        chunks = [chunk async for chunk in client.send_message(request)]
        assert chunks, "Expected at least one streaming chunk"
        assert "echo:stream me" in _collect_text(chunks)
    finally:
        await client.close()


@pytest.mark.asyncio
async def test_multiturn(asgi_client: httpx.AsyncClient) -> None:
    """Second message carries the task_id/context_id from the first turn."""
    resolver = A2ACardResolver(httpx_client=asgi_client, base_url=_AGENT_BASE)
    card = await resolver.get_agent_card()

    config = ClientConfig(streaming=False, httpx_client=asgi_client)
    client = await create_client(agent=card, client_config=config)
    try:
        # --- first turn ---
        first_request = SendMessageRequest(
            message=new_text_message("turn one", role=Role.ROLE_USER)
        )
        task_id = context_id = None
        async for chunk in client.send_message(first_request):
            if chunk.HasField("task"):
                task_id = chunk.task.id
                context_id = chunk.task.context_id

        assert task_id, "First turn must return a task id"
        assert context_id, "First turn must return a context id"

        # --- second turn (same task / context) ---
        second_request = SendMessageRequest(
            message=new_text_message(
                "turn two",
                role=Role.ROLE_USER,
                task_id=task_id,
                context_id=context_id,
            )
        )
        chunks = [chunk async for chunk in client.send_message(second_request)]
        assert chunks, "Second turn must return chunks"
        assert "echo:turn two" in _collect_text(chunks)
    finally:
        await client.close()
