"""
BrokerAuthClient — high-level API for agent-side broker authentication.

Ties together keypair generation, registration, DPoP proof creation, and nonce handling.
"""

from __future__ import annotations

import requests

from broker_auth.dpop import create_dpop_proof
from broker_auth.keypair import generate_es256_keypair
from broker_auth.registration import confirm_key, register_key


class BrokerAuthClient:
    """
    High-level client for MCP Broker Auth Profile.

    Usage:
        client = BrokerAuthClient("http://broker", "gateway-token")
        client.enroll()  # Register + confirm ES256 key
        response = client.make_request("POST", "http://broker/api/v1/mcp/invoke/", json={"tool": "x"})
    """

    def __init__(self, broker_url, gateway_token, keystore_path=None):
        self.broker_url = broker_url
        self.gateway_token = gateway_token
        self._private_key = None
        self._public_jwk = None
        self._kid = None
        self._enrolled = False
        self._keystore = None

        if keystore_path:
            from broker_auth.keystore import KeyStore

            self._keystore = KeyStore(keystore_path)
            if self._keystore.exists():
                data = self._keystore.load()
                if data:
                    self._private_key = data["private_key"]
                    self._public_jwk = data.get("public_jwk")
                    self._kid = data["kid"]
                    self._enrolled = True

    @property
    def is_enrolled(self):
        """Whether the client has a confirmed key."""
        return self._enrolled

    def enroll(self):
        """
        Full registration + confirmation flow.

        Generates an ES256 keypair, registers with the broker, and confirms
        via challenge-response.
        """
        self._private_key, self._public_jwk = generate_es256_keypair()

        # Register
        reg_data = register_key(self.broker_url, self.gateway_token, self._public_jwk)
        self._kid = reg_data["bot_key_id"]

        # Confirm
        confirmed = confirm_key(
            self.broker_url,
            self.gateway_token,
            self._kid,
            reg_data["challenge"],
            self._private_key,
            self._public_jwk,
        )

        if confirmed:
            self._enrolled = True
            if self._keystore:
                self._keystore.save(self._private_key, self._kid, self.broker_url, self._public_jwk)

    def make_request(self, method, url, nonce=None, **kwargs):
        """
        Make an authenticated request with DPoP proof (if enrolled).

        Handles nonce retry: if the server returns 401 with a DPoP-Nonce header,
        retries with the provided nonce.
        """
        headers = kwargs.pop("headers", {})
        headers["Authorization"] = f"Bearer {self.gateway_token}"

        if self._enrolled and self._private_key and self._kid:
            dpop_proof = create_dpop_proof(self._private_key, self._kid, method, url, nonce=nonce)
            headers["DPoP"] = dpop_proof

        response = requests.request(method, url, headers=headers, timeout=30, **kwargs)

        # Handle nonce retry
        if response.status_code == 401 and "DPoP-Nonce" in response.headers:
            server_nonce = response.headers["DPoP-Nonce"]
            headers_retry = dict(headers)
            if self._enrolled and self._private_key and self._kid:
                dpop_proof = create_dpop_proof(self._private_key, self._kid, method, url, nonce=server_nonce)
                headers_retry["DPoP"] = dpop_proof
            response = requests.request(method, url, headers=headers_retry, timeout=30, **kwargs)

        return response

    def rotate_key(self):
        """Generate new keypair, register and confirm with broker."""
        from broker_auth.registration import rotate_key as _rotate

        new_private_key, new_public_jwk = generate_es256_keypair()
        rot_data = _rotate(
            self.broker_url,
            self.gateway_token,
            self._private_key,
            self._public_jwk,
            new_public_jwk,
        )

        # Confirm new key
        confirmed = confirm_key(
            self.broker_url,
            self.gateway_token,
            rot_data["bot_key_id"],
            rot_data["challenge"],
            new_private_key,
            new_public_jwk,
        )

        if confirmed:
            self._private_key = new_private_key
            self._public_jwk = new_public_jwk
            self._kid = rot_data["bot_key_id"]
            if self._keystore:
                self._keystore.save(self._private_key, self._kid, self.broker_url, self._public_jwk)
