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

import asyncio
from collections.abc import Awaitable, Callable, Iterable
from dataclasses import dataclass
from typing import Any
from uuid import uuid4

import grpc.aio
import httpx
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 (
    PushNotificationConfigStore,
    TaskStore,
    TaskUpdater,
)
from a2a.types import Part, Task, TaskState, TaskStatus

from ag2 import Agent
from ag2.a2a import A2AConfig, A2AServer, build_card
from ag2.a2a.testing import (
    ToolRoundsConfig,
    make_test_client_factory,
    make_test_rest_client_factory,
    pick_free_port,
)
from ag2.events import (
    BaseEvent,
    ModelMessageChunk,
)
from ag2.testing import TestConfig, TrackingConfig, Turn


def a2a_config(agent: Agent) -> A2AConfig:
    """The agent's ``A2AConfig``.

    ``Agent.config`` is declared ``ModelConfig | None``, but the task and
    push-notification admin helpers take an ``A2AConfig`` — narrow once here
    instead of handing them an unchecked union at every call site.
    """
    config = agent.config
    assert isinstance(config, A2AConfig), f"{agent.name} is not backed by an A2AConfig"
    return config


@dataclass(slots=True)
class A2APair:
    server: A2AServer
    server_agent: Agent
    client: Agent
    tracking: TrackingConfig


@dataclass(slots=True)
class ExecutorPair:
    server: A2AServer
    executor: A2AAgentExecutorBase
    client: Agent


@dataclass(slots=True)
class GrpcPair:
    server: A2AServer
    server_agent: Agent
    client: Agent
    tracking: TrackingConfig
    grpc_url: str
    grpc_server: grpc.aio.Server


class PromptThenAckExecutor(A2AAgentExecutorBase):
    def __init__(self, prompt: str) -> None:
        self._prompt = prompt
        self.received_user_text: str | None = None

    async def execute(self, request_context: RequestContext, event_queue: EventQueue) -> None:
        msg = request_context.message
        if msg is None:
            return
        task_id = msg.task_id or uuid4().hex
        context_id = msg.context_id or uuid4().hex
        updater = TaskUpdater(event_queue, task_id, context_id)

        if request_context.current_task is None:
            await event_queue.enqueue_event(
                Task(
                    id=task_id,
                    context_id=context_id,
                    status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
                ),
            )
            await updater.start_work()
            await updater.requires_input(
                message=updater.new_agent_message(parts=[Part(text=self._prompt)]),
            )
            return

        text = "".join(p.text for p in msg.parts if p.text)
        self.received_user_text = text
        await updater.complete(
            message=updater.new_agent_message(parts=[Part(text=f"echo: {text}")]),
        )

    async def cancel(self, request_context: RequestContext, event_queue: EventQueue) -> None:
        task = request_context.current_task
        if task is None:
            return
        updater = TaskUpdater(event_queue, task.id, task.context_id)
        await updater.cancel()


class Switchboard(httpx.AsyncBaseTransport):
    """Send each request to whichever replica is `active` — a load balancer that changes its mind."""

    def __init__(self, *servers: A2AServer, url: str) -> None:
        self._transports = [httpx.ASGITransport(app=server.build_jsonrpc(url=url)) for server in servers]
        self.active = 0

    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
        return await self._transports[self.active].handle_async_request(request)


class GatedExecutor(A2AAgentExecutorBase):
    """Starts a task, then holds it `working` until the test opens the gate."""

    def __init__(self) -> None:
        self.gate = asyncio.Event()
        self.working = asyncio.Event()
        self.task_id = ""

    async def execute(self, request_context: RequestContext, event_queue: EventQueue) -> None:
        task_id, context_id = request_context.task_id or uuid4().hex, request_context.context_id or uuid4().hex
        self.task_id = task_id
        updater = TaskUpdater(event_queue, task_id, context_id)
        await event_queue.enqueue_event(
            Task(id=task_id, context_id=context_id, status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED))
        )
        await updater.start_work()
        self.working.set()
        await self.gate.wait()
        await updater.complete(message=updater.new_agent_message(parts=[Part(text="finished")]))

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


def _tool_rounds(initial: Turn, after_tool: "Turn | None") -> ToolRoundsConfig:
    """The server answers with ``initial``, then with ``after_tool`` once the tool results come back."""
    return ToolRoundsConfig(initial, after_tool if after_tool is not None else "")


def make_pair(
    initial: Turn,
    after_tool: "Turn | None" = None,
    *,
    server_tools: Iterable[Callable[..., object]] = (),
    client_tools: Iterable[Callable[..., object]] = (),
    server_url: str = "http://test",
    streaming: bool = True,
    task_store: TaskStore | None = None,
    push_config_store: PushNotificationConfigStore | None = None,
) -> A2APair:
    tracking = TrackingConfig(_tool_rounds(initial, after_tool))
    server_agent = Agent("server-agent", config=tracking)
    for tool in server_tools:
        server_agent.tool(tool)

    server_kwargs: dict[str, Any] = {}
    if task_store is not None:
        server_kwargs["task_store"] = task_store
    if push_config_store is not None:
        server_kwargs["push_config_store"] = push_config_store

    server = A2AServer(server_agent, **server_kwargs)
    factory = make_test_client_factory(server, url=server_url)

    client_config = A2AConfig(
        card_url=server_url,
        httpx_client_factory=factory,
        streaming=streaming,
    )
    client_agent = Agent("client-agent", config=client_config)
    for tool in client_tools:
        client_agent.tool(tool)

    return A2APair(server=server, server_agent=server_agent, client=client_agent, tracking=tracking)


def make_executor_pair(
    executor: A2AAgentExecutorBase,
    *,
    server_url: str = "http://test",
    streaming: bool = False,
    task_store: TaskStore | None = None,
    push_config_store: PushNotificationConfigStore | None = None,
    hitl_hook: Callable[..., Any] | None = None,
) -> ExecutorPair:
    server_agent = Agent("server-stub", config=TestConfig("unused"))

    server_kwargs: dict[str, Any] = {}
    if task_store is not None:
        server_kwargs["task_store"] = task_store
    if push_config_store is not None:
        server_kwargs["push_config_store"] = push_config_store

    server = A2AServer(server_agent, executor=executor, **server_kwargs)
    factory = make_test_client_factory(server, url=server_url)

    client_kwargs: dict[str, Any] = {}
    if hitl_hook is not None:
        client_kwargs["hitl_hook"] = hitl_hook

    client = Agent(
        "client",
        config=A2AConfig(card_url=server_url, httpx_client_factory=factory, streaming=streaming),
        **client_kwargs,
    )
    return ExecutorPair(server=server, executor=executor, client=client)


def make_rest_pair(
    initial: Turn,
    after_tool: "Turn | None" = None,
    *,
    server_url: str = "http://test",
    streaming: bool = False,
) -> A2APair:
    tracking = TrackingConfig(_tool_rounds(initial, after_tool))
    server_agent = Agent("server-agent", config=tracking)
    server = A2AServer(server_agent)
    factory = make_test_rest_client_factory(server, url=server_url)

    client = Agent(
        "client-agent",
        config=A2AConfig(
            card_url=server_url,
            httpx_client_factory=factory,
            prefer="rest",
            streaming=streaming,
        ),
    )
    return A2APair(server=server, server_agent=server_agent, client=client, tracking=tracking)


async def start_grpc_pair(
    initial: Turn,
    after_tool: "Turn | None" = None,
    *,
    host: str = "127.0.0.1",
    streaming: bool = False,
) -> GrpcPair:
    tracking = TrackingConfig(_tool_rounds(initial, after_tool))
    server_agent = Agent("server-agent", config=tracking)
    server = A2AServer(server_agent)

    grpc_url = f"{host}:{pick_free_port(host)}"
    card = build_card(server_agent, url=grpc_url, transports=("grpc",), grpc_url=grpc_url)
    grpc_server = server.build_grpc(bind=grpc_url, grpc_url=grpc_url, card=card)
    await grpc_server.start()

    client = Agent(
        "client-agent",
        config=A2AConfig(
            card_url=grpc_url,
            preset_card=card,
            prefer="grpc",
            streaming=streaming,
        ),
    )
    return GrpcPair(
        server=server,
        server_agent=server_agent,
        client=client,
        tracking=tracking,
        grpc_url=grpc_url,
        grpc_server=grpc_server,
    )


_Script = Callable[[TaskUpdater, asyncio.Event], Awaitable[None]]


class ScriptedExecutor(A2AAgentExecutorBase):
    """Runs a fresh task through ``script``; ``release`` is opened by the test, never by a timer."""

    def __init__(self, script: _Script) -> None:
        self._script = script
        self.release = asyncio.Event()

    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 self._script(updater, self.release)

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


class EventCollector:
    """Subscriber that records the events it is given; ``on_event`` is handed to ``stream.subscribe``."""

    def __init__(self) -> None:
        self.events: list[BaseEvent] = []

    async def on_event(self, event: BaseEvent) -> None:
        self.events.append(event)

    @property
    def chunks(self) -> list[str]:
        return [e.content for e in self.events if isinstance(e, ModelMessageChunk)]


_Turn = Callable[[TaskUpdater], Awaitable[None]]


async def _complete_with_ok(updater: TaskUpdater) -> None:
    await updater.complete(message=updater.new_agent_message([Part(text="ok")]))


class ToolRoundTripExecutor(A2AAgentExecutorBase):
    """Runs ``first_turn`` for a new task and ``continuation`` once the client answers its request for input."""

    def __init__(self, first_turn: _Turn, continuation: _Turn = _complete_with_ok) -> None:
        self._first_turn = first_turn
        self._continuation = continuation

    async def execute(self, request_context: RequestContext, event_queue: EventQueue) -> None:
        msg = request_context.message
        assert msg is not None
        resumed = request_context.current_task is not None
        task_id = msg.task_id or uuid4().hex
        context_id = msg.context_id or uuid4().hex
        if not resumed:
            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 (self._continuation if resumed else self._first_turn)(updater)

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