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

Ties together keypair generation, registration, DPoP proof creation, and nonce handling.
Supports step-up authentication via OAuth 2.0 Device Authorization Grant.
"""

from __future__ import annotations

import logging
import time

logger = logging.getLogger("broker_auth")

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, discover_endpoints, register_key


class StepUpRequiredError(Exception):
    """Raised when the broker returns 403 with error='step_up_required'.

    Attributes:
        trigger: The step-up trigger (e.g. 'mfa', 'consent').
        description: Human-readable description from the broker.
        response_data: Full JSON response body.
    """

    def __init__(self, trigger: str, description: str, response_data: dict):
        self.trigger = trigger
        self.description = description
        self.response_data = response_data
        super().__init__(f"Step-up required ({trigger}): {description}")


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/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
        self._endpoints = 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

    @property
    def endpoints(self):
        """Discovered broker endpoints (cached after first call)."""
        if self._endpoints is None:
            self._endpoints = discover_endpoints(self.broker_url)
        return self._endpoints

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

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

        # Register
        reg_data = register_key(
            self.broker_url,
            self.gateway_token,
            self._public_jwk,
            register_path=self.endpoints["key_register"],
        )
        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,
            confirm_path=self.endpoints["key_confirm"],
        )

        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, access_token=self.gateway_token
            )
            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, access_token=self.gateway_token
                )
                headers_retry["DPoP"] = dpop_proof
            response = requests.request(method, url, headers=headers_retry, timeout=30, **kwargs)

        # Handle step-up authentication requirement
        if response.status_code == 403:
            try:
                body = response.json()
            except (ValueError, AttributeError):
                body = {}
            if body.get("error") == "step_up_required":
                raise StepUpRequiredError(
                    trigger=body.get("trigger", "unknown"),
                    description=body.get("description", "Step-up authentication required"),
                    response_data=body,
                )

        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,
            rotate_path=self.endpoints["key_rotate"],
        )

        # 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,
            confirm_path=self.endpoints["key_confirm"],
        )

        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)

    # ------------------------------------------------------------------
    # Device Authorization Grant (step-up authentication)
    # ------------------------------------------------------------------

    def request_device_authorization(self) -> dict:
        """
        Initiate the OAuth 2.0 Device Authorization Grant.

        POST {broker_url}/auth/device/authorize

        Returns:
            dict with device_code, user_code, verification_uri,
            verification_uri_complete, expires_in, interval.
        """
        url = f"{self.broker_url.rstrip('/')}/auth/device/authorize"
        headers = {"Authorization": f"Bearer {self.gateway_token}"}
        response = requests.post(url, headers=headers, timeout=30)
        response.raise_for_status()
        return response.json()

    def poll_device_token(self, device_code: str, interval: int = 5, timeout: int = 1800) -> dict:
        """
        Poll the device token endpoint until the user approves, denies, or it expires.

        POST {broker_url}/auth/device/token

        Args:
            device_code: The device_code from request_device_authorization().
            interval: Seconds between polls (default 5, may be overridden by server).
            timeout: Maximum seconds to wait before giving up (default 1800).

        Returns:
            dict with access_token (or error details) on completion.

        Raises:
            TimeoutError: If polling exceeds timeout.
            RuntimeError: If the authorization is denied or expired by the server.
        """
        url = f"{self.broker_url.rstrip('/')}/auth/device/token"
        headers = {"Authorization": f"Bearer {self.gateway_token}"}
        deadline = time.monotonic() + timeout
        poll_interval = interval

        while time.monotonic() < deadline:
            response = requests.post(
                url,
                headers=headers,
                json={"device_code": device_code},
                timeout=30,
            )

            if response.status_code == 200:
                return response.json()

            try:
                body = response.json()
            except (ValueError, AttributeError):
                body = {}

            error = body.get("error", "")

            if error == "authorization_pending":
                time.sleep(poll_interval)
                continue
            elif error == "slow_down":
                poll_interval += 5
                time.sleep(poll_interval)
                continue
            elif error in ("access_denied", "expired_token"):
                raise RuntimeError(f"Device authorization failed: {error} — {body.get('error_description', '')}")
            else:
                raise RuntimeError(f"Unexpected device token error: {error} — {body.get('error_description', '')}")

        raise TimeoutError(f"Device authorization polling timed out after {timeout}s")

    def handle_step_up(self) -> bool:
        """
        Orchestrate the full device-authorization step-up flow.

        1. Requests device authorization from the broker.
        2. Prints user instructions to stderr.
        3. Polls until approved or denied.

        Returns:
            True if step-up succeeded and a new token was obtained.
            False if denied, expired, or timed out.
        """
        try:
            auth_data = self.request_device_authorization()
        except requests.HTTPError:
            return False

        # Display instructions to the user
        user_code = auth_data.get("user_code", "N/A")
        verification_uri = auth_data.get("verification_uri_complete") or auth_data.get("verification_uri", "N/A")
        logger.warning(
            "[Step-Up Auth] Visit: %s  |  Enter code: %s",
            verification_uri,
            user_code,
        )

        interval = auth_data.get("interval", 5)
        expires_in = auth_data.get("expires_in", 1800)

        try:
            result = self.poll_device_token(
                auth_data["device_code"],
                interval=interval,
                timeout=expires_in,
            )
            # If the broker returns a new access token, update gateway_token
            if "access_token" in result:
                self.gateway_token = result["access_token"]
            return True
        except (RuntimeError, TimeoutError):
            return False
