"""Ed25519 attestation signing and verification."""

from __future__ import annotations

import base64
from pathlib import Path
from typing import Any

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.serialization import (
    load_pem_private_key,
    load_pem_public_key,
)

from complyform.core.errors import (
    AttestationKeyNotFoundError,
    AttestationVerificationFailedError,
)

PUBLIC_KEY_PATH = Path(__file__).parent.parent / "auth" / "ed25519_public.key"


def load_public_key() -> Ed25519PublicKey:
    """Load the bundled Ed25519 public key."""
    if not PUBLIC_KEY_PATH.exists():
        raise AttestationKeyNotFoundError
    key_bytes = PUBLIC_KEY_PATH.read_bytes()
    key = load_pem_public_key(key_bytes)
    if not isinstance(key, Ed25519PublicKey):
        msg = "Public key is not Ed25519"
        raise AttestationVerificationFailedError(msg)
    return key


def sign_payload(payload_bytes: bytes, private_key_pem: bytes) -> bytes:
    """Sign a payload with an Ed25519 private key."""
    key = load_pem_private_key(private_key_pem, password=None)
    if not isinstance(key, Ed25519PrivateKey):
        msg = "Private key is not Ed25519"
        raise TypeError(msg)
    return key.sign(payload_bytes)


def verify_signature(
    payload_bytes: bytes,
    signature: bytes,
    public_key: Ed25519PublicKey,
) -> bool:
    """Verify an Ed25519 signature. Returns True if valid, False otherwise."""
    try:
        public_key.verify(signature, payload_bytes)
    except InvalidSignature:
        return False
    return True


def build_dsse_envelope(
    statement_json: str,
    private_key_pem: bytes | None = None,
) -> dict[str, Any]:
    """Build a DSSE envelope, optionally signing with Ed25519.

    Args:
        statement_json: Canonical JSON string of the in-toto statement.
        private_key_pem: PEM-encoded Ed25519 private key, or None for unsigned.

    Returns:
        DSSE envelope dict with payloadType, payload, and signatures.
    """
    if isinstance(statement_json, str):
        payload_bytes = statement_json.encode("utf-8")
    else:
        payload_bytes = statement_json
    payload_b64 = base64.b64encode(payload_bytes).decode("ascii")

    # DSSE PAE (Pre-Authentication Encoding)
    pae = _dsse_pae("application/vnd.in-toto+json", payload_bytes)

    signatures: list[dict[str, str]] = []
    if private_key_pem is not None:
        sig_bytes = sign_payload(pae, private_key_pem)
        signatures.append(
            {
                "keyid": "",
                "sig": base64.b64encode(sig_bytes).decode("ascii"),
            }
        )

    return {
        "payloadType": "application/vnd.in-toto+json",
        "payload": payload_b64,
        "signatures": signatures,
    }


def _dsse_pae(payload_type: str, payload: bytes) -> bytes:
    """DSSE Pre-Authentication Encoding (PAE).

    PAE(type, body) = "DSSEv1" + SP + len(type) + SP + type + SP + len(body) + SP + body
    """
    type_bytes = payload_type.encode("utf-8")
    return (
        b"DSSEv1 "
        + str(len(type_bytes)).encode("ascii")
        + b" "
        + type_bytes
        + b" "
        + str(len(payload)).encode("ascii")
        + b" "
        + payload
    )
