#!/usr/bin/env python3
"""
Federation state (off-repo, leaf module)
========================================
The federation receiver's persistence + audit store. Two jobs, both off-repo so they
never sync to a shared remote:

  1. IDEMPOTENCY — which inbound peer-message Message-IDs have already been applied, so a
     re-poll never double-files the same envelope.
  2. AUDIT — an append-only JSONL trail of every decision (who sent it, what intents were
     applied vs. routed to review vs. rejected, and why). The human-readable counterpart
     to the git history of the vault edits the receiver makes.

It also tracks outbound Message-IDs (per peer address) the EGRESS path sent, so a peer's
reply can be proven to thread off a genuine message from this vault — symmetric with the
inbound threading gate.

No dependencies of its own (so every consumer imports it without a cycle) and entirely
defensive: a missing/corrupt store degrades to "nothing processed yet" rather than
raising. The store path defaults OUTSIDE the vault checkout (~/.alexandria/) and is
overridable via FEDERATION_STATE / FEDERATION_AUDIT — it holds Message-IDs and raw email
excerpts that must never be committed to the vault's git history.
"""
import json
import os
import tempfile
import time
from contextlib import contextmanager
from pathlib import Path

try:
    import fcntl
except ImportError:                                    # pragma: no cover — non-POSIX
    fcntl = None

# Default OUTSIDE the vault (never committed): a per-user state dir. Both paths are env
# overridable so a deployment can point them at its own off-repo secrets location.
_DEFAULT_DIR = Path(os.environ.get("ALEXANDRIA_STATE_DIR",
                                   str(Path.home() / ".alexandria")))
STATE_PATH = Path(os.environ.get(
    "FEDERATION_STATE", str(_DEFAULT_DIR / "federation-state.json")))
AUDIT_PATH = Path(os.environ.get(
    "FEDERATION_AUDIT", str(_DEFAULT_DIR / "federation-audit.jsonl")))
_KEEP_IDS = 1000   # cap each rolling id list so the state file can't grow without bound
# A v2 envelope older than the relay's 14-day blob TTL can only be a malicious relay
# re-injecting a captured blob — the freshness gate (receive.py) rejects it, so the
# dedupe set only needs to remember ids within this window (time-evicted below).
MAX_ENVELOPE_AGE = 14 * 24 * 3600


@contextmanager
def _locked():
    """A cross-process exclusive lock around a load-modify-save, so the served worker
    thread and a separate `alexandria federation tick` process never lose each other's
    write (both used to read-modify-write the whole JSON unlocked). Best-effort:
    degrades to no lock where fcntl is unavailable."""
    if fcntl is None:
        yield
        return
    try:
        STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
        fh = open(str(STATE_PATH) + ".lock", "w")
    except OSError:
        yield
        return
    try:
        fcntl.flock(fh, fcntl.LOCK_EX)
        yield
    finally:
        try:
            fcntl.flock(fh, fcntl.LOCK_UN)
        except OSError:
            pass
        fh.close()


def load():
    try:
        return json.loads(STATE_PATH.read_text(encoding="utf-8"))
    except Exception:
        return {}


def save(state):
    """Atomic write (temp + replace) so a concurrent reader never sees a torn file."""
    try:
        STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
        fd, tmp = tempfile.mkstemp(dir=str(STATE_PATH.parent), prefix=".fedstate.")
        try:
            os.write(fd, json.dumps(state, indent=2).encode("utf-8"))
            os.fsync(fd)
            os.close(fd)
            os.replace(tmp, str(STATE_PATH))
        except BaseException:
            try:
                os.close(fd)
            except OSError:
                pass
            try:
                os.unlink(tmp)
            except FileNotFoundError:
                pass
            raise
    except Exception:
        pass


def _update(mutate):
    """Atomically load → mutate(state) → save under the cross-process lock."""
    with _locked():
        s = load()
        mutate(s)
        save(s)


def _push(seq, value):
    seq = [x for x in (seq or []) if x != value]
    seq.append(value)
    return seq[-_KEEP_IDS:]


# ---- inbound idempotency ---------------------------------------------------
# One dedupe implementation, keyed by which rolling list it consults, so the two
# transports never collide in the same store and share exactly one degrade-to-
# empty / cap machinery (a third list-name in P5 is then a one-liner).
def _is_seen(key, value):
    return bool(value) and value in set(load().get(key, []))


def _mark_seen(key, value):
    if not value:
        return
    _update(lambda s: s.__setitem__(key, _push(s.get(key, []), value)))


def is_processed(msg_id):
    """True if this inbound peer-message Message-ID (V1 email) was already handled."""
    return _is_seen("processed_ids", msg_id)


def mark_processed(msg_id):
    _mark_seen("processed_ids", msg_id)


# V2 relay envelopes dedupe by envelope id, keyed PER PEER and TIME-EVICTED. A V2
# envelope's id is authenticated on open (envelope binds it into the sealed header),
# so a successful open guarantees the id is the sender's. Per-peer so a flood from one
# peer can't evict another's ids; time-evicted (MAX_ENVELOPE_AGE) rather than a fixed
# global window so a malicious relay can't replay an old applied envelope by rolling
# its id out of a bounded list — the freshness gate keeps stale envelopes out entirely.
def is_processed_envelope(peer_fp, eid):
    """True if this V2 envelope id from this peer was already applied/handled."""
    m = load().get("processed_envelope_ids", {})
    peer = m.get(peer_fp) if isinstance(m, dict) else None
    return bool(eid) and isinstance(peer, dict) and eid in peer


def mark_processed_envelope(peer_fp, eid, *, now=None):
    """Record this V2 envelope id as processed — an atomic test-and-set. Returns True
    when the id was NEWLY recorded (this caller owns processing it) and False when it
    was already present (someone else filed it — skip). Presence is decided INSIDE the
    cross-process lock, so two concurrent pollers can never both see "new": a separate
    is_processed_envelope() read followed by a mark is exactly the check-then-act race
    this closes. An id-less envelope can't be tracked: nothing is recorded and True is
    returned (process it — the pre-existing behavior for empty ids)."""
    if not eid:
        return True
    now = time.time() if now is None else now
    cutoff = now - MAX_ENVELOPE_AGE
    newly = {"recorded": True}

    def mut(s):
        m = s.get("processed_envelope_ids")
        # Migrate the old flat-list format (pre-fix) by starting fresh — a one-time
        # reset that at worst re-files a handful of very recent envelopes.
        if not isinstance(m, dict):
            m = {}
        peer = m.get(peer_fp)
        if isinstance(peer, dict) and eid in peer:
            # Already recorded (matches what is_processed_envelope would report):
            # leave the entry untouched — the old skip path never refreshed it either.
            newly["recorded"] = False
            return
        peer = {k: v for k, v in peer.items()
                if isinstance(v, (int, float)) and v >= cutoff} if isinstance(peer, dict) else {}
        peer[eid] = now
        m[peer_fp] = peer
        s["processed_envelope_ids"] = m

    _update(mut)
    return newly["recorded"]


# ---- per-peer intent-envelope counter (review_first_n, §5.5) -----------------
def peer_intent_count(fp):
    """How many non-receipt envelopes we've accepted from this peer so far."""
    counts = load().get("peer_intent_counts", {})
    return int(counts.get(fp, 0)) if isinstance(counts, dict) else 0


def bump_peer_intent_count(fp):
    """Increment and return this peer's non-receipt envelope count. The returned value
    is the one THIS call wrote, captured inside the locked mutation — a separate
    unlocked re-read could return another writer's later value."""
    bumped = {"count": 0}

    def mut(s):
        counts = s.get("peer_intent_counts")
        if not isinstance(counts, dict):
            counts = {}
        counts[fp] = bumped["count"] = int(counts.get(fp, 0)) + 1
        s["peer_intent_counts"] = counts

    _update(mut)
    return bumped["count"]


def clear_peer(fp):
    """Forget all fingerprint-keyed state for a removed peer — its dedupe window AND
    its review counter — so a later re-pair of the same brain starts fresh (the
    review_first_n new-peer buffer applies again)."""
    def mut(s):
        for key in ("processed_envelope_ids", "peer_intent_counts"):
            d = s.get(key)
            if isinstance(d, dict):
                d.pop(fp, None)
    _update(mut)


# ---- outbound threading (egress → peer reply) ------------------------------
def record_sent_message(to_addr, msg_id):
    """Remember a Message-ID this vault sent TO a peer's ingress address, so a reply FROM
    that address can be proven to thread off a genuine message we sent. Best-effort.
    Goes through _update like every other writer: a bare load→save here would hold no
    lock and could clobber a concurrent locked writer's update (atomic replace prevents
    torn reads, not lost updates)."""
    if not to_addr or not msg_id:
        return
    a = to_addr.strip().lower()

    def mut(s):
        m = s.get("sent_message_ids")
        if not isinstance(m, dict):
            m = {}
        m[a] = _push(m.get(a, []), msg_id)
        s["sent_message_ids"] = m

    _update(mut)


def sent_message_ids(*to_addrs):
    """All Message-IDs we sent to any of `to_addrs` (case-insensitive union)."""
    m = load().get("sent_message_ids", {})
    if not isinstance(m, dict):
        return []
    out = []
    for a in to_addrs:
        out += m.get((a or "").strip().lower(), [])
    return out


# ---- audit trail -----------------------------------------------------------
_AUDIT_MAX_BYTES = 5 * 1024 * 1024   # rotate once the JSONL grows past ~5 MB


def append_audit(record):
    """Append one decision record (provenance + what was applied / reviewed / rejected) to
    the off-repo JSONL audit trail. Best-effort; never raises into the caller. The file
    is rotated (kept + .1) once it passes _AUDIT_MAX_BYTES so it can't grow without
    bound on an always-on brain."""
    try:
        AUDIT_PATH.parent.mkdir(parents=True, exist_ok=True)
        try:
            if AUDIT_PATH.stat().st_size > _AUDIT_MAX_BYTES:
                os.replace(str(AUDIT_PATH), str(AUDIT_PATH) + ".1")
        except FileNotFoundError:
            pass
        with open(AUDIT_PATH, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
    except Exception:
        pass
