"""
Key registration and confirmation with the broker.
"""

from __future__ import annotations

import base64
import json

import requests
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec

from broker_auth.keypair import compute_jwk_thumbprint


def _sign_challenge(private_key, challenge_data):
    """Sign challenge data with ES256 private key, return base64url signature."""
    signature = private_key.sign(challenge_data.encode(), ec.ECDSA(hashes.SHA256()))
    return base64.urlsafe_b64encode(signature).rstrip(b"=").decode()


def discover_endpoints(broker_url):
    """
    Fetch broker endpoints from the discovery document.

    GET {broker_url}/.well-known/mcp-broker

    Returns: dict of endpoint paths (key_register, key_confirm, key_rotate, invoke, tools, etc.)
    """
    response = requests.get(
        f"{broker_url}/.well-known/mcp-broker",
        timeout=30,
    )
    response.raise_for_status()
    return response.json()["endpoints"]


def register_key(broker_url, gateway_token, public_jwk, register_path=None):
    """
    Register an ES256 public key with the broker.

    Args:
        broker_url: Base URL of the broker
        gateway_token: Bearer token for authentication
        public_jwk: ES256 public key in JWK format
        register_path: Optional endpoint path (discovered automatically if not provided)

    Returns: {"bot_key_id": ..., "challenge": ..., "expires_in_seconds": ...}
    """
    if not register_path:
        endpoints = discover_endpoints(broker_url)
        register_path = endpoints["key_register"]

    response = requests.post(
        f"{broker_url}{register_path}",
        json={"public_jwk": public_jwk},
        headers={"Authorization": f"Bearer {gateway_token}"},
        timeout=30,
    )
    response.raise_for_status()
    return response.json()


def confirm_key(broker_url, gateway_token, bot_key_id, challenge, private_key, public_jwk, confirm_path=None):
    """
    Confirm key registration via challenge-response.

    Signs "{challenge}.{jwk_thumbprint}" with the private key.

    Returns: True if confirmed.
    """
    if not confirm_path:
        endpoints = discover_endpoints(broker_url)
        confirm_path = endpoints["key_confirm"]

    thumbprint = compute_jwk_thumbprint(public_jwk)
    challenge_data = f"{challenge}.{thumbprint}"
    signature = _sign_challenge(private_key, challenge_data)

    response = requests.post(
        f"{broker_url}{confirm_path}",
        json={
            "bot_key_id": bot_key_id,
            "challenge": challenge,
            "signature": signature,
        },
        headers={"Authorization": f"Bearer {gateway_token}"},
        timeout=30,
    )
    response.raise_for_status()
    return response.json().get("key_confirmed", False)


def rotate_key(broker_url, gateway_token, old_private_key, old_public_jwk, new_public_jwk, rotate_path=None):
    """
    Rotate bot key: sign rotation request with current key, provide new key.

    Returns: {"bot_key_id": ..., "challenge": ...}
    """
    if not rotate_path:
        endpoints = discover_endpoints(broker_url)
        rotate_path = endpoints["key_rotate"]

    rotation_data = json.dumps({"new_public_jwk": new_public_jwk}, sort_keys=True, separators=(",", ":"))
    signature = _sign_challenge(old_private_key, rotation_data)

    response = requests.post(
        f"{broker_url}{rotate_path}",
        json={"new_public_jwk": new_public_jwk, "signature": signature},
        headers={"Authorization": f"Bearer {gateway_token}"},
        timeout=30,
    )
    response.raise_for_status()
    return response.json()
