"""Client-side SDK surface. App teams use RemoteAgentClient + send_data()
instead of touching a2a.client / httpx / auth directly.
"""
from __future__ import annotations

from typing import Any

import httpx
from a2a.client import Client, ClientConfig, create_client
from a2a.helpers import get_data_parts, new_data_message
from a2a.types.a2a_pb2 import Role, SendMessageRequest

from a2a_passport.auth.base import TokenProvider
from a2a_passport.config import settings
from a2a_passport.retry import default_retry


class _BearerAuth(httpx.Auth):
    """Attaches a token from any TokenProvider, scoped to one audience."""

    def __init__(self, token_provider: TokenProvider, audience: str):
        self._token_provider = token_provider
        self._audience = audience

    def sync_auth_flow(self, request: httpx.Request):  # pragma: no cover
        raise RuntimeError("this auth flow is async-only; use httpx.AsyncClient")

    async def async_auth_flow(self, request: httpx.Request):
        token = await self._token_provider.get_token(self._audience)
        request.headers["Authorization"] = f"Bearer {token}"
        yield request


class RemoteAgentClient:
    """A configured, authenticated, retrying handle to one remote agent.

    `token_provider` is any object satisfying the TokenProvider protocol --
    pass a2a_passport.auth.oauth2.OAuth2ClientCredentialsProvider for the
    common case (client-credentials against your IdP), or a custom
    implementation for anything else. One RemoteAgentClient targets one
    audience; construct a separate one per remote agent you call.
    """

    def __init__(
        self,
        base_url: str,
        token_provider: TokenProvider,
        *,
        audience: str,
        request_timeout_seconds: float | None = None,
    ):
        self._base_url = base_url
        self._auth = _BearerAuth(token_provider, audience=audience)
        self._timeout = request_timeout_seconds or settings.request_timeout_seconds
        self._httpx_client: httpx.AsyncClient | None = None
        self._client: Client | None = None

    async def __aenter__(self) -> RemoteAgentClient:
        self._httpx_client = httpx.AsyncClient(auth=self._auth, timeout=self._timeout)
        config = ClientConfig(streaming=False, httpx_client=self._httpx_client)
        self._client = await create_client(agent=self._base_url, client_config=config)
        return self

    async def __aexit__(self, *exc):
        if self._client:
            await self._client.close()
        if self._httpx_client:
            await self._httpx_client.aclose()

    @default_retry
    async def send_data(self, data: dict[str, Any]) -> dict[str, Any]:
        """Send a structured data message; return the first artifact's data dict.

        Raises RuntimeError if the task did not complete successfully.
        """
        assert self._client is not None, "use 'async with RemoteAgentClient(...)'"
        msg = new_data_message(data, role=Role.ROLE_USER)
        req = SendMessageRequest(message=msg)
        last_task = None
        async for chunk in self._client.send_message(req):
            last_task = chunk.task

        if last_task is None:
            raise RuntimeError("no task returned from remote agent")

        from a2a.types.a2a_pb2 import TaskState  # local import: keeps pb2 out of module surface

        if last_task.status.state != TaskState.TASK_STATE_COMPLETED:
            raise RuntimeError(f"remote task did not complete: state={last_task.status.state}")
        if not last_task.artifacts:
            return {}
        results = get_data_parts(last_task.artifacts[0].parts)
        return results[0] if results else {}
