"""Async client for the UCM Agent Commerce Protocol."""

import time
import uuid
from typing import Optional

import httpx

from .crypto import canonical_json, sha256_hex, sign_message
from .exceptions import raise_for_acp_error


class UCMClient:
    """Async client for the UCM Agent Commerce Protocol.

    Supports two authentication modes:
    - **Ed25519 signatures**: Pass ``private_key`` and ``public_key`` to the constructor.
    - **API Key (Bearer token)**: Pass ``api_key`` (e.g. ``ucm_key_...``).

    You may provide both; Ed25519 is preferred for endpoints that require it
    (publish, authorize, delete), while the API key is used as a fallback for
    buy/balance/history.
    """

    def __init__(
        self,
        registry_url: str,
        private_key: str | None = None,
        public_key: str | None = None,
        api_key: str | None = None,
        timeout: float = 30.0,
    ):
        self.registry_url = registry_url.rstrip('/')
        self.private_key = private_key
        self.public_key = public_key
        self.api_key = api_key
        self._client = httpx.AsyncClient(timeout=timeout)

        if not private_key and not api_key:
            raise ValueError('Either (private_key, public_key) or api_key must be provided')
        if private_key and not public_key:
            raise ValueError('public_key is required when private_key is provided')

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _sign_headers(self, method: str, path: str, body: str = '') -> dict[str, str]:
        """Generate Ed25519 signed authentication headers."""
        if not self.private_key or not self.public_key:
            raise ValueError('Ed25519 keys are required for signed requests')

        timestamp = int(time.time() * 1000)
        nonce = str(uuid.uuid4())

        payload = canonical_json({
            'body_hash': sha256_hex(body),
            'method': method,
            'nonce': nonce,
            'path': path,
            'timestamp': timestamp,
        })

        signature = sign_message(self.private_key, sha256_hex(payload))

        return {
            'X-ACP-Pubkey': self.public_key,
            'X-ACP-Timestamp': str(timestamp),
            'X-ACP-Nonce': nonce,
            'X-ACP-Signature': signature,
            'Content-Type': 'application/json',
        }

    def _bearer_headers(self) -> dict[str, str]:
        """Generate Bearer token authentication headers."""
        if not self.api_key:
            raise ValueError('api_key is required for Bearer auth')
        return {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json',
        }

    def _auth_headers(self, method: str, path: str, body: str = '') -> dict[str, str]:
        """Choose the best available auth method (Ed25519 preferred, then API key)."""
        if self.private_key:
            return self._sign_headers(method, path, body)
        return self._bearer_headers()

    async def _request(self, method: str, path: str, body: str | None = None,
                       headers: dict[str, str] | None = None) -> dict:
        """Send a request and handle ACP error responses."""
        url = f'{self.registry_url}{path}'
        kwargs: dict = {'headers': headers or {}}
        if body is not None:
            kwargs['content'] = body

        resp = await self._client.request(method, url, **kwargs)

        if resp.status_code >= 400:
            try:
                data = resp.json()
            except Exception:
                resp.raise_for_status()
                return {}
            raise_for_acp_error(resp.status_code, data)

        return resp.json() if resp.content else {}

    # ------------------------------------------------------------------
    # Public endpoints (no auth required)
    # ------------------------------------------------------------------

    async def get_onboarding_md(self) -> str:
        """Fetch the agent onboarding guide (onboarding.md) as markdown text."""
        resp = await self._client.get(f'{self.registry_url}/onboarding.md')
        resp.raise_for_status()
        return resp.text

    async def discover(self, need: str, tags: Optional[list[str]] = None,
                       limit: int = 5) -> dict:
        """Discover services matching a natural language need.

        Args:
            need: Natural language description of what you need.
            tags: Optional list of tags to filter by.
            limit: Maximum number of results (default 5, max 20).

        Returns:
            Dict with ``services`` list.
        """
        body_obj: dict = {'need': need, 'limit': min(limit, 20)}
        if tags:
            body_obj['tags'] = tags
        body = canonical_json(body_obj)

        return await self._request('POST', '/v1/discover', body,
                                   headers={'Content-Type': 'application/json'})

    async def service_info(self, service_id: str) -> dict:
        """Get detailed information about a specific service.

        Args:
            service_id: The service identifier (e.g. ``acme/web-search``).

        Returns:
            Full ServiceDescriptor dict.
        """
        return await self._request('GET', f'/v1/services/{service_id}')

    async def verify_token(self, token: str) -> dict:
        """Verify an access token's validity.

        Args:
            token: The ``acp_sk_...`` access token.

        Returns:
            Dict with token payload if valid.
        """
        body = canonical_json({'token': token})
        return await self._request('POST', '/v1/token/verify', body,
                                   headers={'Content-Type': 'application/json'})

    async def deposit_info(self) -> dict:
        """Get Solana escrow deposit information.

        Returns:
            Dict with escrow address, cluster, supported tokens, and balances.
        """
        return await self._request('GET', '/v1/deposit/info')

    # ------------------------------------------------------------------
    # Agent management (no auth or API key)
    # ------------------------------------------------------------------

    async def register(self, name: str, description: str | None = None) -> dict:
        """Register a new agent and obtain an API key.

        Args:
            name: Agent display name (max 255 chars).
            description: Optional agent description.

        Returns:
            Dict with ``agent_id``, ``api_key``, ``name``, ``budget``, ``expires_at``.
        """
        body_obj: dict = {'name': name}
        if description:
            body_obj['description'] = description
        body = canonical_json(body_obj)

        result = await self._request('POST', '/v1/agents/register', body,
                                     headers={'Content-Type': 'application/json'})

        # Auto-configure api_key if not already set
        if not self.api_key and 'api_key' in result:
            self.api_key = result['api_key']

        return result

    async def bind_wallet(self, wallet: str, signature: str) -> dict:
        """Bind a Solana wallet to the agent (required for paid services).

        Args:
            wallet: Solana wallet address (base58).
            signature: Ed25519 signature of ``ucm:bind-wallet:{agent_id}``.

        Returns:
            Dict with ``agent_id``, ``wallet``, confirmation message.
        """
        path = '/v1/agents/bind-wallet'
        body = canonical_json({'wallet': wallet, 'signature': signature})
        headers = self._auth_headers('POST', path, body)
        return await self._request('POST', path, body, headers)

    async def recover(self, wallet: str, timestamp: int, signature: str) -> dict:
        """Recover agent account via wallet signature (issues new API key).

        Args:
            wallet: Solana wallet address (base58).
            timestamp: Current time in milliseconds.
            signature: Ed25519 signature of ``ucm:recover:{agent_id}:{timestamp}``.

        Returns:
            Dict with ``agent_id``, ``api_key``, confirmation message.
        """
        body = canonical_json({
            'wallet': wallet,
            'timestamp': timestamp,
            'signature': signature,
        })
        result = await self._request('POST', '/v1/agents/recover', body,
                                     headers={'Content-Type': 'application/json'})

        # Auto-update api_key after recovery
        if 'api_key' in result:
            self.api_key = result['api_key']

        return result

    # ------------------------------------------------------------------
    # Authenticated endpoints (Ed25519 or API key)
    # ------------------------------------------------------------------

    async def buy(self, service_id: str, plan: str = 'per-call') -> dict:
        """Purchase access to a service.

        Args:
            service_id: The service identifier (e.g. ``acme/web-search``).
            plan: Pricing plan (``per-call``, ``per-unit``, ``subscription``).

        Returns:
            Dict with ``tx_id``, ``access_token``, ``service_endpoint``,
            ``token_expires_at``, ``amount_charged``, ``budget_remaining``.
        """
        path = '/v1/buy'
        body = canonical_json({'service_id': service_id, 'plan': plan})
        headers = self._auth_headers('POST', path, body)
        return await self._request('POST', path, body, headers)

    async def balance(self) -> dict:
        """Check agent's budget balance and usage.

        Returns:
            Dict with budget limits and current spending counters.
        """
        path = '/v1/balance'
        headers = self._auth_headers('GET', path)
        return await self._request('GET', path, headers=headers)

    async def history(self, limit: int = 20) -> dict:
        """Get transaction history for the agent.

        Args:
            limit: Maximum number of receipts (1–100, default 20).

        Returns:
            Dict with ``receipts`` list.
        """
        path = f'/v1/history?limit={min(max(limit, 1), 100)}'
        headers = self._auth_headers('GET', path)
        return await self._request('GET', path, headers=headers)

    # ------------------------------------------------------------------
    # Ed25519-only endpoints (provider/operator operations)
    # ------------------------------------------------------------------

    async def publish(self, descriptor: dict) -> dict:
        """Publish a new service to the marketplace (Ed25519 signed).

        Args:
            descriptor: Full ServiceDescriptor dict with ``acp_version``,
                ``service_id``, ``name``, ``endpoints``, ``provider``,
                ``base_url``, ``pricing``, etc.

        Returns:
            Dict with ``service_id``, ``status``, ``descriptor_hash``.
        """
        path = '/v1/publish'
        body = canonical_json(descriptor)
        headers = self._sign_headers('POST', path, body)
        return await self._request('POST', path, body, headers)

    async def update_service(self, service_id: str, update: dict) -> dict:
        """Update an existing service (Ed25519 signed).

        Args:
            service_id: The service identifier.
            update: Partial ServiceDescriptor with fields to update.

        Returns:
            Dict with ``service_id``, ``status``, ``descriptor_hash``.
        """
        path = f'/v1/publish/{service_id}'
        body = canonical_json(update)
        headers = self._sign_headers('PUT', path, body)
        return await self._request('PUT', path, body, headers)

    async def delete_service(self, service_id: str) -> dict:
        """Delete (delist) a service from the marketplace (Ed25519 signed).

        Args:
            service_id: The service identifier.

        Returns:
            Dict with ``service_id`` and ``status: delisted``.
        """
        path = f'/v1/publish/{service_id}'
        headers = self._sign_headers('DELETE', path)
        return await self._request('DELETE', path, headers=headers)

    async def authorize(self, agent: str, rules: dict,
                        operator: str | None = None) -> dict:
        """Create a budget authorization for an agent (Ed25519 signed by operator).

        Args:
            agent: Agent pubkey or ``solana:<pubkey>``.
            rules: Budget rules dict with ``max_per_tx``, ``max_daily``,
                ``max_monthly``, ``currency``, ``expires_at``.
            operator: Operator identifier. Defaults to ``solana:<public_key>``.

        Returns:
            Dict with ``agent``, ``status: authorized``, ``rules``.
        """
        path = '/v1/authorize'
        auth_obj = {
            'acp_version': '0.1',
            'type': 'budget_authorization',
            'agent': agent,
            'operator': operator or f'solana:{self.public_key}',
            'rules': rules,
            'created_at': time.strftime('%Y-%m-%dT%H:%M:%S.000Z', time.gmtime()),
        }
        body = canonical_json(auth_obj)
        headers = self._sign_headers('POST', path, body)
        return await self._request('POST', path, body, headers)

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def close(self):
        """Close the HTTP client."""
        await self._client.aclose()

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        await self.close()
