"""
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 register_key(broker_url, gateway_token, public_jwk):
    """
    Register an ES256 public key with the broker.

    POST {broker_url}/api/v1/bots/keys/register/

    Returns: {"bot_key_id": ..., "challenge": ..., "expires_in_seconds": ...}
    """
    response = requests.post(
        f"{broker_url}/api/v1/bots/keys/register/",
        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 key registration via challenge-response.

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

    Returns: True if confirmed.
    """
    thumbprint = compute_jwk_thumbprint(public_jwk)
    challenge_data = f"{challenge}.{thumbprint}"
    signature = _sign_challenge(private_key, challenge_data)

    response = requests.post(
        f"{broker_url}/api/v1/bots/keys/confirm/",
        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 bot key: sign rotation request with current key, provide new key.

    Returns: {"bot_key_id": ..., "challenge": ...}
    """
    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}/api/v1/bots/keys/rotate/",
        json={"new_public_jwk": new_public_jwk, "signature": signature},
        headers={"Authorization": f"Bearer {gateway_token}"},
        timeout=30,
    )
    response.raise_for_status()
    return response.json()
