"""Backwards compatibility layer for legacy A2A clients."""

import warnings

from collections.abc import AsyncGenerator
from typing import Any

import httpx

from pygeai_proxy.vendor.a2a.client.errors import A2AClientJSONRPCError
from pygeai_proxy.vendor.a2a.client.middleware import ClientCallContext, ClientCallInterceptor
from pygeai_proxy.vendor.a2a.client.transports.jsonrpc import JsonRpcTransport
from pygeai_proxy.vendor.a2a.types import (
    AgentCard,
    CancelTaskRequest,
    CancelTaskResponse,
    CancelTaskSuccessResponse,
    GetTaskPushNotificationConfigParams,
    GetTaskPushNotificationConfigRequest,
    GetTaskPushNotificationConfigResponse,
    GetTaskPushNotificationConfigSuccessResponse,
    GetTaskRequest,
    GetTaskResponse,
    GetTaskSuccessResponse,
    JSONRPCErrorResponse,
    SendMessageRequest,
    SendMessageResponse,
    SendMessageSuccessResponse,
    SendStreamingMessageRequest,
    SendStreamingMessageResponse,
    SendStreamingMessageSuccessResponse,
    SetTaskPushNotificationConfigRequest,
    SetTaskPushNotificationConfigResponse,
    SetTaskPushNotificationConfigSuccessResponse,
    TaskIdParams,
    TaskResubscriptionRequest,
)


class A2AClient:
    """[DEPRECATED] Backwards compatibility wrapper for the JSON-RPC client."""

    def __init__(
        self,
        httpx_client: httpx.AsyncClient,
        agent_card: AgentCard | None = None,
        url: str | None = None,
        interceptors: list[ClientCallInterceptor] | None = None,
    ):
        warnings.warn(
            'A2AClient is deprecated and will be removed in a future version. '
            'Use ClientFactory to create a client with a JSON-RPC transport.',
            DeprecationWarning,
            stacklevel=2,
        )
        self._transport = JsonRpcTransport(
            httpx_client, agent_card, url, interceptors
        )

    async def send_message(
        self,
        request: SendMessageRequest,
        *,
        http_kwargs: dict[str, Any] | None = None,
        context: ClientCallContext | None = None,
    ) -> SendMessageResponse:
        """Sends a non-streaming message request to the agent.

        Args:
            request: The `SendMessageRequest` object containing the message and configuration.
            http_kwargs: Optional dictionary of keyword arguments to pass to the
                underlying httpx.post request.
            context: The client call context.

        Returns:
            A `SendMessageResponse` object containing the agent's response (Task or Message) or an error.

        Raises:
            A2AClientHTTPError: If an HTTP error occurs during the request.
            A2AClientJSONError: If the response body cannot be decoded as JSON or validated.
        """
        if not context and http_kwargs:
            context = ClientCallContext(state={'http_kwargs': http_kwargs})

        try:
            result = await self._transport.send_message(
                request.params, context=context
            )
            return SendMessageResponse(
                root=SendMessageSuccessResponse(
                    id=request.id, jsonrpc='2.0', result=result
                )
            )
        except A2AClientJSONRPCError as e:
            return SendMessageResponse(JSONRPCErrorResponse(error=e.error))

    async def send_message_streaming(
        self,
        request: SendStreamingMessageRequest,
        *,
        http_kwargs: dict[str, Any] | None = None,
        context: ClientCallContext | None = None,
    ) -> AsyncGenerator[SendStreamingMessageResponse, None]:
        """Sends a streaming message request to the agent and yields responses as they arrive.

        This method uses Server-Sent Events (SSE) to receive a stream of updates from the agent.

        Args:
            request: The `SendStreamingMessageRequest` object containing the message and configuration.
            http_kwargs: Optional dictionary of keyword arguments to pass to the
                underlying httpx.post request. A default `timeout=None` is set but can be overridden.
            context: The client call context.

        Yields:
            `SendStreamingMessageResponse` objects as they are received in the SSE stream.
            These can be Task, Message, TaskStatusUpdateEvent, or TaskArtifactUpdateEvent.

        Raises:
            A2AClientHTTPError: If an HTTP or SSE protocol error occurs during the request.
            A2AClientJSONError: If an SSE event data cannot be decoded as JSON or validated.
        """
        if not context and http_kwargs:
            context = ClientCallContext(state={'http_kwargs': http_kwargs})

        async for result in self._transport.send_message_streaming(
            request.params, context=context
        ):
            yield SendStreamingMessageResponse(
                root=SendStreamingMessageSuccessResponse(
                    id=request.id, jsonrpc='2.0', result=result
                )
            )

    async def get_task(
        self,
        request: GetTaskRequest,
        *,
        http_kwargs: dict[str, Any] | None = None,
        context: ClientCallContext | None = None,
    ) -> GetTaskResponse:
        """Retrieves the current state and history of a specific task.

        Args:
            request: The `GetTaskRequest` object specifying the task ID and history length.
            http_kwargs: Optional dictionary of keyword arguments to pass to the
                underlying httpx.post request.
            context: The client call context.

        Returns:
            A `GetTaskResponse` object containing the Task or an error.

        Raises:
            A2AClientHTTPError: If an HTTP error occurs during the request.
            A2AClientJSONError: If the response body cannot be decoded as JSON or validated.
        """
        if not context and http_kwargs:
            context = ClientCallContext(state={'http_kwargs': http_kwargs})
        try:
            result = await self._transport.get_task(
                request.params, context=context
            )
            return GetTaskResponse(
                root=GetTaskSuccessResponse(
                    id=request.id, jsonrpc='2.0', result=result
                )
            )
        except A2AClientJSONRPCError as e:
            return GetTaskResponse(root=JSONRPCErrorResponse(error=e.error))

    async def cancel_task(
        self,
        request: CancelTaskRequest,
        *,
        http_kwargs: dict[str, Any] | None = None,
        context: ClientCallContext | None = None,
    ) -> CancelTaskResponse:
        """Requests the agent to cancel a specific task.

        Args:
            request: The `CancelTaskRequest` object specifying the task ID.
            http_kwargs: Optional dictionary of keyword arguments to pass to the
                underlying httpx.post request.
            context: The client call context.

        Returns:
            A `CancelTaskResponse` object containing the updated Task with canceled status or an error.

        Raises:
            A2AClientHTTPError: If an HTTP error occurs during the request.
            A2AClientJSONError: If the response body cannot be decoded as JSON or validated.
        """
        if not context and http_kwargs:
            context = ClientCallContext(state={'http_kwargs': http_kwargs})
        try:
            result = await self._transport.cancel_task(
                request.params, context=context
            )
            return CancelTaskResponse(
                root=CancelTaskSuccessResponse(
                    id=request.id, jsonrpc='2.0', result=result
                )
            )
        except A2AClientJSONRPCError as e:
            return CancelTaskResponse(JSONRPCErrorResponse(error=e.error))

    async def set_task_callback(
        self,
        request: SetTaskPushNotificationConfigRequest,
        *,
        http_kwargs: dict[str, Any] | None = None,
        context: ClientCallContext | None = None,
    ) -> SetTaskPushNotificationConfigResponse:
        """Sets or updates the push notification configuration for a specific task.

        Args:
            request: The `SetTaskPushNotificationConfigRequest` object specifying the task ID and configuration.
            http_kwargs: Optional dictionary of keyword arguments to pass to the
                underlying httpx.post request.
            context: The client call context.

        Returns:
            A `SetTaskPushNotificationConfigResponse` object containing the confirmation or an error.

        Raises:
            A2AClientHTTPError: If an HTTP error occurs during the request.
            A2AClientJSONError: If the response body cannot be decoded as JSON or validated.
        """
        if not context and http_kwargs:
            context = ClientCallContext(state={'http_kwargs': http_kwargs})
        try:
            result = await self._transport.set_task_callback(
                request.params, context=context
            )
            return SetTaskPushNotificationConfigResponse(
                root=SetTaskPushNotificationConfigSuccessResponse(
                    id=request.id, jsonrpc='2.0', result=result
                )
            )
        except A2AClientJSONRPCError as e:
            return SetTaskPushNotificationConfigResponse(
                JSONRPCErrorResponse(error=e.error)
            )

    async def get_task_callback(
        self,
        request: GetTaskPushNotificationConfigRequest,
        *,
        http_kwargs: dict[str, Any] | None = None,
        context: ClientCallContext | None = None,
    ) -> GetTaskPushNotificationConfigResponse:
        """Retrieves the push notification configuration for a specific task.

        Args:
            request: The `GetTaskPushNotificationConfigRequest` object specifying the task ID.
            http_kwargs: Optional dictionary of keyword arguments to pass to the
                underlying httpx.post request.
            context: The client call context.

        Returns:
            A `GetTaskPushNotificationConfigResponse` object containing the configuration or an error.

        Raises:
            A2AClientHTTPError: If an HTTP error occurs during the request.
            A2AClientJSONError: If the response body cannot be decoded as JSON or validated.
        """
        if not context and http_kwargs:
            context = ClientCallContext(state={'http_kwargs': http_kwargs})
        params = request.params
        if isinstance(params, TaskIdParams):
            params = GetTaskPushNotificationConfigParams(id=request.params.id)
        try:
            result = await self._transport.get_task_callback(
                params, context=context
            )
            return GetTaskPushNotificationConfigResponse(
                root=GetTaskPushNotificationConfigSuccessResponse(
                    id=request.id, jsonrpc='2.0', result=result
                )
            )
        except A2AClientJSONRPCError as e:
            return GetTaskPushNotificationConfigResponse(
                JSONRPCErrorResponse(error=e.error)
            )

    async def resubscribe(
        self,
        request: TaskResubscriptionRequest,
        *,
        http_kwargs: dict[str, Any] | None = None,
        context: ClientCallContext | None = None,
    ) -> AsyncGenerator[SendStreamingMessageResponse, None]:
        """Reconnects to get task updates.

        This method uses Server-Sent Events (SSE) to receive a stream of updates from the agent.

        Args:
            request: The `TaskResubscriptionRequest` object containing the task information to reconnect to.
            http_kwargs: Optional dictionary of keyword arguments to pass to the
                underlying httpx.post request. A default `timeout=None` is set but can be overridden.
            context: The client call context.

        Yields:
            `SendStreamingMessageResponse` objects as they are received in the SSE stream.
            These can be Task, Message, TaskStatusUpdateEvent, or TaskArtifactUpdateEvent.

        Raises:
            A2AClientHTTPError: If an HTTP or SSE protocol error occurs during the request.
            A2AClientJSONError: If an SSE event data cannot be decoded as JSON or validated.
        """
        if not context and http_kwargs:
            context = ClientCallContext(state={'http_kwargs': http_kwargs})

        async for result in self._transport.resubscribe(
            request.params, context=context
        ):
            yield SendStreamingMessageResponse(
                root=SendStreamingMessageSuccessResponse(
                    id=request.id, jsonrpc='2.0', result=result
                )
            )

    async def get_card(
        self,
        *,
        http_kwargs: dict[str, Any] | None = None,
        context: ClientCallContext | None = None,
    ) -> AgentCard:
        """Retrieves the authenticated card (if necessary) or the public one.

        Args:
            http_kwargs: Optional dictionary of keyword arguments to pass to the
                underlying httpx.post request.
            context: The client call context.

        Returns:
            A `AgentCard` object containing the card or an error.

        Raises:
            A2AClientHTTPError: If an HTTP error occurs during the request.
            A2AClientJSONError: If the response body cannot be decoded as JSON or validated.
        """
        if not context and http_kwargs:
            context = ClientCallContext(state={'http_kwargs': http_kwargs})
        return await self._transport.get_card(context=context)
