"""Brain identity: one X25519 keypair, its fingerprint, and verification words.

The fingerprint (``ax1…``) is the brain's federation address — peers allowlist
it, the relay names mailboxes by it, envelopes carry it in ``from``/``to``.
It is a pure function of the public key, so the private key file IS the
identity: **never regenerate over an existing key** — a rotated key silently
orphans every peer relationship (their allowlist still names the old
fingerprint, and nothing they send can be opened).

Storage (docs/V4-FEDERATION.md §2.1): ``$ALEXANDRIA_STATE_DIR/federation/``
holds ``brain.key`` (32 raw private-key bytes, mode 0600) and ``brain.pub``
(32 raw public-key bytes). State-dir paths are read lazily — never at import —
so tests and deployments can repoint ``ALEXANDRIA_STATE_DIR`` per process
(same contract as engine/federation_state.py, minus its import-time wart).
"""
import base64
import hashlib
import os
import re
import tempfile
from pathlib import Path

from alexandria_vault.federation import sodium
from alexandria_vault.federation.words import WORDS256

FINGERPRINT_PREFIX = "ax1"
# The canonical fingerprint format, owned here and reused by envelope.parse so the
# validator can never drift from what fingerprint() actually emits. Fully anchored
# (``\A``/``\Z``, not ``$``) and exact-length: a real fingerprint is ``ax1`` + 32
# base32 chars = 35 chars, so a trailing newline or an over-length relay-injected
# string is rejected, not accepted.
FINGERPRINT_RE = re.compile(r"\Aax1[a-z2-7]{32}\Z")
# Domain separator for verification words (frozen — part of the wire contract).
_VERIFY_DOMAIN = b"axon-verify-v1"
KEY_FILE = "brain.key"
PUB_FILE = "brain.pub"


def is_fingerprint(value):
    """True iff ``value`` is a well-formed ``ax1…`` fingerprint (str, exact shape)."""
    return isinstance(value, str) and FINGERPRINT_RE.match(value) is not None


def _default_home():
    """``$ALEXANDRIA_STATE_DIR/federation`` (default ``~/.alexandria/federation``).

    Read lazily, never at import, and ``expanduser``'d so a literal ``~`` in the
    env var resolves to the home dir rather than a stray ``./~`` in the CWD.
    """
    root = Path(os.environ.get("ALEXANDRIA_STATE_DIR")
                or str(Path.home() / ".alexandria")).expanduser()
    return root / "federation"


def _ensure_home(home):
    """Create the state home if absent and assert 0700 — for BOTH the default and
    an explicitly-passed path, so a caller-supplied secrets dir is never left
    world-listable (the perms discipline used to apply only on the default path)."""
    home.mkdir(parents=True, exist_ok=True)
    os.chmod(home, 0o700)
    return home


def federation_home(create=False):
    """The state home — the ONLY place private keys, mailbox bearers, and deposit
    tokens ever live. If a code path would put one of those in config, the vault,
    a log line, or a commit — stop, that's the bug (§8 of the build plan)."""
    home = _default_home()
    return _ensure_home(home) if create else home


def _write_secret_atomically(path, data):
    """Write a secret file atomically at 0600 — no world-readable window, no
    torn/truncated file, and first-writer-wins under a concurrent first run.

    A temp file in the same directory is fully written + fsync'd, then hard-linked
    into place (``os.link`` is atomic and fails closed if the target already
    exists). So a crash mid-write can only leave the temp file (swept below),
    never a partial ``brain.key`` — the state the length guard would otherwise
    refuse forever. A lost link race means another process created the real key
    first; the caller then simply loads that one. The repo has other bespoke
    secure-file writers (scripts/cli.py, scripts/claude_auth.py); P4/P5's bearer
    and deposit-token writers should share this one rather than add a fourth.
    """
    fd, tmp = tempfile.mkstemp(dir=str(path.parent), prefix="." + path.name + ".")
    try:
        os.write(fd, data)
        os.fsync(fd)
    finally:
        os.close(fd)
    os.chmod(tmp, 0o600)  # mkstemp is already 0600; assert it under any umask
    try:
        os.link(tmp, str(path))  # atomic + exclusive: FileExistsError if we raced
    except FileExistsError:
        pass                     # another process created the key first — fine
    finally:
        try:
            os.unlink(tmp)
        except FileNotFoundError:
            pass


def load_or_create(home=None):
    """Return ``(PrivateKey, PublicKey)``, generating exactly once.

    Idempotent: an existing ``brain.key`` is always loaded, never replaced.
    ``brain.pub`` is (re)derived from the private key if missing or stale, so
    the pair on disk can't drift apart.
    """
    nacl = sodium()
    home = _ensure_home(Path(home) if home is not None else _default_home())
    key_path = home / KEY_FILE
    if not key_path.exists():
        # Generate + atomically install. If we lose the race, the winner's key is
        # already at key_path and we read it back below — both processes converge
        # on one identity instead of the loser crashing on O_EXCL.
        _write_secret_atomically(key_path, bytes(nacl.public.PrivateKey.generate()))
    raw = key_path.read_bytes()
    if len(raw) != 32:
        raise ValueError(
            f"{key_path} is corrupt ({len(raw)} bytes, expected 32) — refusing to "
            "touch it; restore it from backup rather than regenerating (a new key "
            "is a new identity and orphans every paired peer)")
    priv = nacl.public.PrivateKey(raw)
    os.chmod(key_path, 0o600)
    pub = priv.public_key
    pub_path = home / PUB_FILE
    if not pub_path.exists() or pub_path.read_bytes() != bytes(pub):
        pub_path.write_bytes(bytes(pub))
    return priv, pub


def try_load(home=None):
    """``(PrivateKey, PublicKey)`` if a key already exists, else ``(None, None)`` —
    a load-ONLY lookup that never generates. For read-only surfaces (e.g. `peer
    list`) that must not create an identity as a side effect."""
    home = Path(home) if home is not None else _default_home()
    key_path = home / KEY_FILE
    if not key_path.exists():
        return None, None
    return load_or_create(home)


def _pub_bytes(public_key):
    raw = bytes(public_key)
    if len(raw) != 32:
        raise ValueError(f"public key must be 32 bytes, got {len(raw)}")
    return raw


def fingerprint(public_key):
    """The brain's address: ``ax1`` + base32(sha256(pub)[:20]), lowercase, unpadded.

    20 hash bytes encode to exactly 32 base32 chars (no padding), so every
    fingerprint is 35 chars matching ``ax1[a-z2-7]{32}``.
    """
    digest = hashlib.sha256(_pub_bytes(public_key)).digest()[:20]
    b32 = base64.b32encode(digest).decode("ascii").lower().rstrip("=")
    return FINGERPRINT_PREFIX + b32


def verification_words(pub_a, pub_b):
    """Three words both owners read aloud to confirm no relay swapped the keys.

    Direction-independent by construction (the two keys are sorted before
    hashing), so each side computes the same three words from its own key and
    the key it *believes* is the peer's — if a middlebox substituted keys, the
    spoken words differ and the humans catch it. Shown at pairing and by
    ``alexandria peers``.

    Strength — read honestly: this is a **24-bit** short-authentication string
    (three bytes of the digest indexing a 256-word list). It reliably catches
    accidental mispairing and an opportunistic/passive relay. It is NOT a
    high-assurance anchor against a *determined active* man-in-the-middle: such
    an attacker who can grind ~2^24 substitute keypairs could find one whose SAS
    collides with the words the two humans expect. Widening the SAS (more words →
    more bits) or adding a key-commitment round is the real fix, but the
    derivation + wordlist are a frozen wire contract, so that is a deliberately
    deferred, wire-versioned change (docs/V4-FEDERATION.md). Do not oversell the
    three words as unconditional proof of identity in owner-facing copy.
    """
    a, b = _pub_bytes(pub_a), _pub_bytes(pub_b)
    lo, hi = (a, b) if a <= b else (b, a)
    d = hashlib.sha256(_VERIFY_DOMAIN + lo + hi).digest()
    return [WORDS256[d[0]], WORDS256[d[1]], WORDS256[d[2]]]
