"""Minimal COSE_Sign1 implementation sufficient for Aliro IssuerAuth.

Reference: RFC 8152 (COSE), \u00a74.4 for the Sig_structure.

COSE_Sign1 wire format:
    [ protected_bstr, unprotected_map, payload_bstr, signature_bstr ]

Sig_structure (what gets signed):
    [ "Signature1", body_protected, external_aad, payload ]
    serialised with deterministic CBOR, then ECDSA-SHA-256.

COSE ECDSA signatures are raw r||s (64 bytes for P-256), not DER.
"""

import cbor2
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.utils import (
    decode_dss_signature,
    encode_dss_signature,
)

COSE_ALG_ES256 = -7
COSE_HDR_ALG = 1
COSE_HDR_KID = 4
_P256_COORD_LEN = 32


class CoseVerificationError(Exception):
    pass


def _sig_structure(protected_bstr: bytes, payload: bytes) -> bytes:
    return cbor2.dumps(["Signature1", protected_bstr, b"", payload], canonical=True)


def _der_to_raw(der_sig: bytes) -> bytes:
    r, s = decode_dss_signature(der_sig)
    return r.to_bytes(_P256_COORD_LEN, "big") + s.to_bytes(_P256_COORD_LEN, "big")


def _raw_to_der(raw_sig: bytes) -> bytes:
    if len(raw_sig) != 2 * _P256_COORD_LEN:
        raise CoseVerificationError(f"expected 64-byte raw signature, got {len(raw_sig)}")
    r = int.from_bytes(raw_sig[:_P256_COORD_LEN], "big")
    s = int.from_bytes(raw_sig[_P256_COORD_LEN:], "big")
    return encode_dss_signature(r, s)


def sign_cose_sign1(
    *,
    private_key: ec.EllipticCurvePrivateKey,
    payload: bytes,
    kid: bytes,
) -> bytes:
    protected_bstr = cbor2.dumps({COSE_HDR_ALG: COSE_ALG_ES256}, canonical=True)
    unprotected = {COSE_HDR_KID: kid}
    to_be_signed = _sig_structure(protected_bstr, payload)
    der_sig = private_key.sign(to_be_signed, ec.ECDSA(hashes.SHA256()))
    raw_sig = _der_to_raw(der_sig)
    return cbor2.dumps([protected_bstr, unprotected, payload, raw_sig], canonical=True)


def verify_cose_sign1(
    signed: bytes,
    public_key: ec.EllipticCurvePublicKey,
) -> bytes:
    try:
        decoded = cbor2.loads(signed)
    except Exception as e:
        raise CoseVerificationError(f"not valid CBOR: {e}") from e

    if not isinstance(decoded, list) or len(decoded) != 4:
        raise CoseVerificationError("COSE_Sign1 must be a 4-element array")
    protected_bstr, _unprotected, payload, raw_sig = decoded
    if not isinstance(protected_bstr, bytes):
        raise CoseVerificationError("protected header must be bstr")
    if not isinstance(payload, bytes):
        raise CoseVerificationError("payload must be bstr")
    if not isinstance(raw_sig, bytes):
        raise CoseVerificationError("signature must be bstr")

    to_be_signed = _sig_structure(protected_bstr, payload)
    der_sig = _raw_to_der(raw_sig)
    try:
        public_key.verify(der_sig, to_be_signed, ec.ECDSA(hashes.SHA256()))
    except InvalidSignature as e:
        raise CoseVerificationError("signature verification failed") from e
    return payload
