# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0

from collections.abc import Callable
from uuid import uuid4

import httpx
import pytest
from a2a.server.agent_execution import AgentExecutor as A2AAgentExecutorBase
from a2a.server.agent_execution import RequestContext
from a2a.server.events import EventQueue
from a2a.server.tasks import InMemoryPushNotificationConfigStore, TaskUpdater
from a2a.types import Task, TaskState, TaskStatus
from fast_depends.pydantic import PydanticSerializer

from ag2 import Agent, Context, MemoryStream
from ag2.a2a import A2AConfig, A2AServer, build_card
from ag2.a2a.client import A2AClient
from ag2.a2a.errors import A2ATaskFailedError
from ag2.a2a.testing import make_test_client_factory
from ag2.events import ModelMessage, ModelMessageChunk, ModelRequest, ModelResponse, TextInput
from ag2.testing import TestConfig

# What an ``Agent`` hands its ``LLMClient``, built the same way.
_SERIALIZER = PydanticSerializer(pydantic_config={"arbitrary_types_allowed": True}, use_fastdepends_errors=False)


class _SpyAsyncClient(httpx.AsyncClient):
    aclose_count = 0

    async def aclose(self) -> None:
        type(self).aclose_count += 1
        await super().aclose()


def _make_spy_factory(server: A2AServer, url: str) -> Callable[[], httpx.AsyncClient]:
    transport = httpx.ASGITransport(app=server.build_jsonrpc(url=url))

    def factory() -> httpx.AsyncClient:
        return _SpyAsyncClient(transport=transport, base_url=url, timeout=30.0)

    return factory


class _AlwaysFailingExecutor(A2AAgentExecutorBase):
    async def execute(self, request_context: RequestContext, event_queue: EventQueue) -> None:
        msg = request_context.message
        assert msg is not None
        task_id = msg.task_id or uuid4().hex
        context_id = msg.context_id or uuid4().hex
        await event_queue.enqueue_event(
            Task(id=task_id, context_id=context_id, status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED)),
        )
        updater = TaskUpdater(event_queue, task_id, context_id)
        await updater.start_work()
        await updater.failed()

    async def cancel(self, request_context: RequestContext, event_queue: EventQueue) -> None:
        return None


@pytest.mark.asyncio
class TestHttpxLifecycle:
    async def test_client_closes_httpx_after_successful_ask(self) -> None:
        _SpyAsyncClient.aclose_count = 0
        server = A2AServer(Agent("server", config=TestConfig("hi")))
        url = "http://test"
        client = Agent("client", config=A2AConfig(card_url=url, httpx_client_factory=_make_spy_factory(server, url)))

        await client.ask("ping")

        assert _SpyAsyncClient.aclose_count == 1

    async def test_client_closes_httpx_when_task_fails(self) -> None:
        _SpyAsyncClient.aclose_count = 0
        server = A2AServer(
            Agent("stub", config=TestConfig("unused")),
            executor=_AlwaysFailingExecutor(),
        )
        url = "http://test"
        client = Agent(
            "client",
            config=A2AConfig(card_url=url, httpx_client_factory=_make_spy_factory(server, url), streaming=False),
        )

        with pytest.raises(A2ATaskFailedError):
            await client.ask("ping")

        assert _SpyAsyncClient.aclose_count == 1

    async def test_closing_leaves_the_client_reusable_for_the_next_turn(self) -> None:
        # ``A2AClient`` is reused across ``reply.ask(...)`` follow-ups and
        # closes its httpx client at the end of every turn, so each turn has
        # to reconnect from scratch rather than reach for a closed client.
        _SpyAsyncClient.aclose_count = 0
        server = A2AServer(Agent("server", config=TestConfig("hi")))
        url = "http://test"
        client = Agent("client", config=A2AConfig(card_url=url, httpx_client_factory=_make_spy_factory(server, url)))

        first = await client.ask("ping")
        second = await first.ask("ping again")

        assert second.body == "hi"
        assert _SpyAsyncClient.aclose_count == 2

    async def test_aclose_on_a_client_that_never_connected_is_a_no_op(self) -> None:
        _SpyAsyncClient.aclose_count = 0
        client = A2AClient(card_url="http://test", httpx_client_factory=_SpyAsyncClient)

        await client.aclose()
        await client.aclose()

        assert _SpyAsyncClient.aclose_count == 0

    async def test_aclose_after_a_turn_does_not_close_the_same_connection_twice(self) -> None:
        # The turn closes in its own ``finally``; an explicit ``aclose()``
        # on top must not reach the httpx client again. Driven directly
        # because ``Agent`` never hands the client back to the caller.
        _SpyAsyncClient.aclose_count = 0
        server = A2AServer(Agent("server", config=TestConfig("hi")))
        url = "http://test"
        client = A2AClient(card_url=url, httpx_client_factory=_make_spy_factory(server, url))

        await client(
            [ModelRequest([TextInput("ping")])],
            Context(stream=MemoryStream()),
            tools=[],
            response_schema=None,
            serializer=_SERIALIZER,
        )
        await client.aclose()

        assert _SpyAsyncClient.aclose_count == 1


@pytest.mark.asyncio
async def test_streamed_chunks_not_duplicated_in_final_message() -> None:
    turn = TestConfig(ModelMessageChunk("he"), ModelMessageChunk("llo"), ModelResponse(message=ModelMessage("hello")))
    server = A2AServer(Agent("server", config=turn))
    url = "http://test"
    factory = make_test_client_factory(server, url=url)
    client = Agent(
        "client",
        config=A2AConfig(card_url=url, httpx_client_factory=factory, streaming=True),
    )

    reply = await client.ask("ping")

    assert reply.response.content == "hello"


@pytest.mark.asyncio
class TestCardImmutability:
    async def test_build_jsonrpc_does_not_mutate_input_card(self) -> None:
        agent = Agent("server", config=TestConfig("hi"))
        card = build_card(agent, url="http://test", transports=("jsonrpc",))
        before_extended = card.capabilities.extended_agent_card
        before_push = card.capabilities.push_notifications

        server = A2AServer(
            agent,
            extended_card=card,
            push_config_store=InMemoryPushNotificationConfigStore(),
        )
        server.build_jsonrpc(url="http://test", card=card)

        assert card.capabilities.extended_agent_card == before_extended
        assert card.capabilities.push_notifications == before_push

    async def test_build_rest_does_not_mutate_input_card(self) -> None:
        agent = Agent("server", config=TestConfig("hi"))
        card = build_card(agent, url="http://test", transports=("rest",))
        before_extended = card.capabilities.extended_agent_card
        before_push = card.capabilities.push_notifications

        server = A2AServer(
            agent,
            extended_card=card,
            push_config_store=InMemoryPushNotificationConfigStore(),
        )
        server.build_rest(url="http://test", card=card)

        assert card.capabilities.extended_agent_card == before_extended
        assert card.capabilities.push_notifications == before_push

    async def test_build_grpc_does_not_mutate_input_card(self) -> None:
        agent = Agent("server", config=TestConfig("hi"))
        card = build_card(agent, url="http://test", transports=("grpc",), grpc_url="localhost:50051")
        before_extended = card.capabilities.extended_agent_card
        before_push = card.capabilities.push_notifications

        server = A2AServer(
            agent,
            extended_card=card,
            push_config_store=InMemoryPushNotificationConfigStore(),
        )
        server.build_grpc(bind="localhost:0", grpc_url="localhost:50051", card=card)

        assert card.capabilities.extended_agent_card == before_extended
        assert card.capabilities.push_notifications == before_push