"""Ed25519 key management and message signing."""

from __future__ import annotations

import base64
import hashlib
import json
import uuid
from pathlib import Path

from nacl.signing import SigningKey, VerifyKey

from .config import KEYS_DIR, PROFILE_FILE, load_json, save_json


def _b64url_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def _b64url_decode(s: str) -> bytes:
    padding = 4 - len(s) % 4
    if padding != 4:
        s += "=" * padding
    return base64.urlsafe_b64decode(s)


def generate_keypair(name: str = "default") -> tuple[str, str, str]:
    """Generate Ed25519 keypair, save to disk, return (did, public_key_b64url, secret_key_path)."""
    sk = SigningKey.generate()
    pk = sk.verify_key

    pk_bytes = bytes(pk)
    did = "did:b2a:z" + _b64url_encode(pk_bytes)[:22]
    pk_b64url = _b64url_encode(pk_bytes)

    key_path = KEYS_DIR / f"{name}.key"
    KEYS_DIR.mkdir(parents=True, exist_ok=True)
    key_path.write_bytes(bytes(sk))
    key_path.chmod(0o600)

    return did, pk_b64url, str(key_path)


def load_signing_key(name: str = "default") -> SigningKey:
    key_path = KEYS_DIR / f"{name}.key"
    if not key_path.exists():
        raise FileNotFoundError(f"No key found at {key_path}. Run 'b2a register' first.")
    return SigningKey(key_path.read_bytes())


def get_public_key_b64url(name: str = "default") -> str:
    sk = load_signing_key(name)
    return _b64url_encode(bytes(sk.verify_key))


def sign_message(
    sender_did: str,
    recipient_did: str,
    payload_bytes: bytes,
    message_id: str | None = None,
    key_name: str = "default",
) -> dict:
    """Create a signed message envelope."""
    sk = load_signing_key(key_name)
    msg_id = message_id or str(uuid.uuid4())

    payload_b64 = _b64url_encode(payload_bytes)

    # Sign: SHA256(message_id || sender_did || recipient_did || payload_bytes)
    to_sign = msg_id.encode() + sender_did.encode() + recipient_did.encode() + payload_bytes
    digest = hashlib.sha256(to_sign).digest()
    signature = sk.sign(digest).signature
    sig_b64 = _b64url_encode(signature)

    return {
        "message_id": msg_id,
        "payload": payload_b64,
        "signature": sig_b64,
        "protocol_version": "0.1",
    }
