from __future__ import annotations

import re
from typing import Any, Dict, Mapping, Union

from nacl.signing import SigningKey

from .canonicalize import canonicalize_for_signing
from .encoding import decode_base64, decode_base64url, encode_base64url, utf8_encode
from .types import KeyPair, to_mapping, to_plain


def _sanitize_pem(pem: str) -> str:
    return re.sub(r"\s+", "", re.sub(r"-----END[^-]+-----", "", re.sub(r"-----BEGIN[^-]+-----", "", pem)))


def generate_key_pair() -> KeyPair:
    signing_key = SigningKey.generate()
    seed = signing_key.encode()
    public_key = signing_key.verify_key.encode()
    return KeyPair(publicKey=public_key, secretKey=seed + public_key)


def load_secret_key(input_key: Union[bytes, bytearray, str]) -> bytes:
    raw: bytes

    if isinstance(input_key, (bytes, bytearray)):
        raw = bytes(input_key)
    elif "-----BEGIN" in input_key:
        raw = decode_base64(_sanitize_pem(input_key))
    else:
        try:
            raw = decode_base64url(input_key)
        except ValueError:
            raw = decode_base64(input_key)

    if len(raw) == 64:
        return raw

    if len(raw) == 32:
        signing_key = SigningKey(raw)
        return raw + signing_key.verify_key.encode()

    raise ValueError("Secret key must be 32-byte seed or 64-byte Ed25519 secret key")


def public_key_from_secret(secret_key: Union[bytes, bytearray, str]) -> bytes:
    secret = load_secret_key(secret_key)
    seed = secret[:32]
    return SigningKey(seed).verify_key.encode()


def sign_receipt(unsigned: Union[Mapping[str, Any], Any], secret_key: Union[bytes, bytearray, str]) -> Dict[str, Any]:
    sk = load_secret_key(secret_key)
    unsigned_obj = to_mapping(unsigned, "unsigned")
    sig_obj = unsigned_obj.get("signature")

    if not isinstance(sig_obj, Mapping):
        raise ValueError("Invalid receipt.signature")

    fallback_public_key = encode_base64url(public_key_from_secret(sk))
    public_key = sig_obj["publicKey"] if "publicKey" in sig_obj and sig_obj["publicKey"] is not None else fallback_public_key

    with_defaults = {
        **unsigned_obj,
        "signature": {
            **to_plain(sig_obj),
            "alg": "Ed25519",
            "canonicalization": "JCS-SORTED-UTF8-NOWS",
            "publicKey": public_key,
        },
    }

    payload = utf8_encode(canonicalize_for_signing(with_defaults))
    signature = SigningKey(sk[:32]).sign(payload).signature

    return {
        **with_defaults,
        "signature": {
            **with_defaults["signature"],
            "sig": encode_base64url(signature),
        },
    }


generateKeyPair = generate_key_pair
loadSecretKey = load_secret_key
publicKeyFromSecret = public_key_from_secret
signReceipt = sign_receipt

__all__ = [
    "generate_key_pair",
    "generateKeyPair",
    "load_secret_key",
    "loadSecretKey",
    "public_key_from_secret",
    "publicKeyFromSecret",
    "sign_receipt",
    "signReceipt",
]
