from __future__ import annotations

from typing import Any, Callable, Dict, Mapping, Optional, Tuple

from nacl.signing import SigningKey, VerifyKey

from api.keys.store__keys import SigningKeyStore, get_default_signing_store
from receipts.cbor import cbor_to_payload, payload_to_cbor
from receipts.cbor_encoder import decode as cbor_decode, encode as cbor_encode
from receipts.models import BaseReceipt, Ed25519Signature, load_receipt

ALG_EDDSA = -8  # COSE alg value for EdDSA on Ed25519
HEADER_PARAM_ALG = 1
HEADER_PARAM_KID = 4

__all__ = [
    "sign_receipt_cose",
    "verify_cose_envelope",
    "load_cose_receipt",
]


def _resolve_signing_record(
    receipt: BaseReceipt,
    *,
    store: SigningKeyStore,
) -> Optional[Tuple[str, SigningKey]]:
    signature_block = getattr(receipt, "signature", None)
    if signature_block is None and hasattr(receipt, "signatures"):
        sigs = getattr(receipt, "signatures") or []
        if sigs:
            signature_block = sigs[0]
    kid: Optional[str] = None
    if isinstance(signature_block, Ed25519Signature):
        kid = signature_block.kid
    elif isinstance(signature_block, Mapping):
        kid_candidate = signature_block.get("kid")
        if isinstance(kid_candidate, str) and kid_candidate:
            kid = kid_candidate
    if kid:
        record = store.get_by_kid(kid)
        if record is None:
            return None
        signing_key = record.signing_key()
        if signing_key is None:
            return None
        return (record.kid, signing_key)

    tenant = receipt.tenant or receipt.domain
    if not tenant:
        return None
    record = store.get(str(tenant))
    if record is None:
        return None
    signing_key = record.signing_key()
    if signing_key is None:
        return None
    return (record.kid, signing_key)


def _build_cose_message(
    *,
    payload_bytes: bytes,
    signing_key: SigningKey,
    kid: str,
) -> bytes:
    protected_map = {HEADER_PARAM_ALG: ALG_EDDSA}
    protected_serialized = cbor_encode(protected_map)
    unprotected_headers = {HEADER_PARAM_KID: kid.encode("utf-8")}

    sig_structure = [
        "Signature1",
        protected_serialized,
        b"",
        payload_bytes,
    ]
    to_sign = cbor_encode(sig_structure)
    signature = signing_key.sign(to_sign).signature

    cose_message = [
        protected_serialized,
        unprotected_headers,
        payload_bytes,
        signature,
    ]
    return cbor_encode(cose_message)


def sign_receipt_cose(
    receipt: BaseReceipt,
    *,
    payload: Optional[Dict[str, Any]] = None,
    store: Optional[SigningKeyStore] = None,
) -> Optional[bytes]:
    """
    Produce a COSE_Sign1 envelope for ``receipt`` using the tenant signing key.

    Returns:
        COSE message bytes when a signing key is available, otherwise ``None``.
    """
    store = store or get_default_signing_store()
    resolved = _resolve_signing_record(receipt, store=store)
    if resolved is None:
        return None
    kid, signing_key = resolved
    payload = dict(payload) if payload is not None else receipt.model_dump(mode="json", by_alias=True)
    # Ensure signature block is embedded in payload for parity with CBOR twin
    sig_block: Optional[Any] = None
    sigs = payload.get("signatures")
    if isinstance(sigs, list) and sigs:
        sig_block = sigs[0]
    if sig_block and isinstance(sig_block, Mapping):
        payload.setdefault("signature", dict(sig_block))
    payload_bytes = payload_to_cbor(payload)
    return _build_cose_message(payload_bytes=payload_bytes, signing_key=signing_key, kid=kid)


def verify_cose_envelope(
    envelope: bytes,
    *,
    key_resolver: Callable[[str], VerifyKey],
) -> Tuple[str, Dict[str, Any]]:
    """
    Verify a COSE_Sign1 envelope and return the decoded payload mapping.

    Args:
        envelope: COSE message bytes.
        key_resolver: Callback returning a ``VerifyKey`` for a given kid.

    Returns:
        Tuple of ``(kid, payload_mapping)`` when verification succeeds.

    Raises:
        RuntimeError: When the envelope structure or signature is invalid.
    """
    try:
        protected_serialized, unprotected_headers, payload_bytes, signature = cbor_decode(envelope)
    except Exception as exc:  # pragma: no cover - defensive guard
        raise RuntimeError("Invalid COSE envelope.") from exc

    if not isinstance(protected_serialized, (bytes, bytearray)):
        raise RuntimeError("COSE protected headers must be bstr.")
    protected_map = {}
    if protected_serialized:
        protected_map = cbor_decode(bytes(protected_serialized))
    if protected_map.get(HEADER_PARAM_ALG) != ALG_EDDSA:
        raise RuntimeError("Unsupported COSE alg header.")

    if not isinstance(unprotected_headers, Mapping):
        raise RuntimeError("COSE unprotected headers must be a map.")
    kid_value = unprotected_headers.get(HEADER_PARAM_KID)
    if not isinstance(kid_value, (bytes, bytearray)) or not kid_value:
        raise RuntimeError("COSE envelope missing kid header.")
    kid = kid_value.decode("utf-8")

    verify_key = key_resolver(kid)

    if not isinstance(payload_bytes, (bytes, bytearray)):
        raise RuntimeError("COSE payload must be bstr.")
    payload_bytes = bytes(payload_bytes)

    sig_structure = [
        "Signature1",
        protected_serialized,
        b"",
        payload_bytes,
    ]
    to_verify = cbor_encode(sig_structure)
    if not isinstance(signature, (bytes, bytearray)):
        raise RuntimeError("COSE signature must be bstr.")
    signature_bytes = bytes(signature)
    verify_key.verify(to_verify, signature_bytes)

    payload = cbor_to_payload(payload_bytes)
    return kid, payload


def load_cose_receipt(
    envelope: bytes,
    *,
    key_resolver: Callable[[str], VerifyKey],
) -> BaseReceipt:
    """
    Verify and decode a COSE envelope into a typed receipt.
    """
    _, payload = verify_cose_envelope(envelope, key_resolver=key_resolver)
    return load_receipt(payload)
