# SPDX-License-Identifier: Apache-2.0
"""Sealed, tiered memory (ADR-0005 + key chains in keys.py).

Each block is distilled into every rung at or below its stamp. Each copy is
encrypted with AES-GCM under the hour key of its (rung, retention class) chain.
Forgetting never touches this store: it advances chain heads, after which the
ciphertext here can no longer be opened.
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field

from cryptography.hazmat.primitives.ciphers.aead import AESGCM

from .keys import Keyring, snap_class


@dataclass
class Block:
    block_id: str
    day: int
    hour: int
    stamp: int                         # rung in force while spoken (0 = off the record)
    texts: dict[int, str]              # distilled text per rung, e.g. {3: ..., 2: ..., 1: ...}
    shorten_days: int | None = None    # "only keep this for a week"


@dataclass
class Store:
    ring: Keyring
    sealed: list[tuple] = field(default_factory=list)   # (block_id, day, hour, rung, cls, nonce, ciphertext)
    gaps: list[tuple] = field(default_factory=list)     # (block_id, day) for off-the-record blocks

    def seal(self, b: Block) -> None:
        if b.stamp == 0:
            self.gaps.append((b.block_id, b.day))       # only the fact that a gap existed
            return
        cls = snap_class(b.shorten_days) if b.shorten_days is not None else "30d"
        for rung in range(1, min(b.stamp, 3) + 1):
            text = b.texts.get(rung)
            if not text:
                continue
            c = "keep" if (rung == 1 and b.shorten_days is None) else cls
            key = self.ring.hour_key(rung, c, b.day, b.hour)[:32]
            nonce = os.urandom(12)
            ct = AESGCM(key).encrypt(nonce, text.encode(), f"{b.block_id}|{rung}".encode())
            self.sealed.append((b.block_id, b.day, b.hour, rung, c, nonce, ct))

    def read(self) -> dict[str, tuple[int, str]]:
        """Best readable rung per block right now."""
        best: dict[str, tuple[int, str]] = {}
        for bid, day, hour, rung, c, nonce, ct in self.sealed:
            try:
                key = self.ring.hour_key(rung, c, day, hour)[:32]
            except KeyError:
                continue                                  # chain advanced: forgotten
            text = AESGCM(key).decrypt(nonce, ct, f"{bid}|{rung}".encode()).decode()
            if bid not in best or rung > best[bid][0]:
                best[bid] = (rung, text)
        return best