"""Ed25519 key management: generation, loading, signing, and fingerprinting.

Mirrors the Node.js bridge-client-node/src/keys.ts implementation exactly,
including file format compatibility (JSON with base64 keys, or raw base64 seed).
"""

from __future__ import annotations

import base64
import hashlib
import json
import os
import stat
from dataclasses import dataclass
from pathlib import Path

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.serialization import (
    Encoding,
    NoEncryption,
    PrivateFormat,
    PublicFormat,
)


@dataclass(frozen=True, slots=True)
class KeyPair:
    public_key: bytes   # 32-byte raw Ed25519 public key
    private_key: bytes  # 32-byte raw Ed25519 seed


def generate_key_pair(save_to: str | Path | None = None) -> KeyPair:
    """Generate a new Ed25519 key pair, optionally persisting to *save_to*."""
    priv = Ed25519PrivateKey.generate()
    raw_private = priv.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())
    raw_public = priv.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)

    pair = KeyPair(public_key=raw_public, private_key=raw_private)

    if save_to is not None:
        p = Path(save_to)
        p.parent.mkdir(parents=True, exist_ok=True)
        data = {
            "publicKey": base64.b64encode(pair.public_key).decode(),
            "privateKey": base64.b64encode(pair.private_key).decode(),
        }
        p.write_text(json.dumps(data, indent=2))
        os.chmod(p, stat.S_IRUSR | stat.S_IWUSR)  # 0o600

    return pair


def load_key_pair(file_path: str | Path) -> KeyPair:
    """Load a key pair from *file_path*, generating one if the file is missing.

    Supports two on-disk formats (same as the Node.js SDK):
      1. JSON ``{"publicKey": "<b64>", "privateKey": "<b64>"}``
      2. Raw base64 private-key seed (32 bytes) — e.g. a K8s Secret mount.
    """
    p = Path(file_path)
    if not p.exists():
        return generate_key_pair(p)

    raw = p.read_text().strip()

    # JSON format
    try:
        data = json.loads(raw)
        if "publicKey" in data and "privateKey" in data:
            return KeyPair(
                public_key=base64.b64decode(data["publicKey"]),
                private_key=base64.b64decode(data["privateKey"]),
            )
    except (json.JSONDecodeError, KeyError):
        pass

    # Raw base64 seed
    seed = base64.b64decode(raw)
    if len(seed) == 32:
        return KeyPair(public_key=public_key_from_private(seed), private_key=seed)

    raise ValueError(f"Invalid key file format at {file_path}")


def fingerprint(public_key: bytes) -> str:
    """SHA-256 hex digest of the 32-byte raw public key (matches Go Bridge)."""
    return hashlib.sha256(public_key).hexdigest()


def sign(message: str, private_key: bytes) -> bytes:
    """Sign *message* (UTF-8) with a 32-byte Ed25519 seed, returning the 64-byte signature."""
    key = Ed25519PrivateKey.from_private_bytes(private_key)
    return key.sign(message.encode())


def public_key_from_private(private_key: bytes) -> bytes:
    """Derive the 32-byte raw public key from a 32-byte Ed25519 seed."""
    key = Ed25519PrivateKey.from_private_bytes(private_key)
    return key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
