# SPDX-License-Identifier: Apache-2.0
"""Forward-secure key chains (v01/research/software-stack.md, "Key management").

One chain per (rung, retention class). The secure element stores only each
chain's head: the seed of the oldest day still readable. Advancing the head
past a day makes that day's key impossible to recompute — that is forgetting.

`SecureElement` is the hardware boundary. `SoftSE` is a stand-in for tests and
the desk rig; the kit uses an ATECC608 implementation of the same interface.
"""
from __future__ import annotations

import hashlib
import hmac
import os
from dataclasses import dataclass
from typing import Protocol

RUNGS = (1, 2, 3)               # R4/R5 copies ride on the R3 chain's expiry
CLASSES = {"1d": 1, "7d": 7, "30d": 30, "keep": None}


class SecureElement(Protocol):
    def read_head(self, slot: str) -> tuple[int, bytes]: ...
    def write_head(self, slot: str, day: int, seed: bytes) -> None: ...
    def bind(self, material: bytes) -> bytes:
        """HMAC with a root secret that never leaves the chip."""
        ...


class SoftSE:
    def __init__(self) -> None:
        self._root = os.urandom(32)
        self._heads: dict[str, tuple[int, bytes]] = {}

    def read_head(self, slot: str) -> tuple[int, bytes]:
        return self._heads[slot]

    def write_head(self, slot: str, day: int, seed: bytes) -> None:
        self._heads[slot] = (day, seed)

    def bind(self, material: bytes) -> bytes:
        return hmac.new(self._root, material, hashlib.sha256).digest()


def snap_class(days: float | None) -> str:
    """'Only keep this for a week' → the nearest class at or below it."""
    if days is None:
        return "keep"
    for name, d in sorted(((k, v) for k, v in CLASSES.items() if v), key=lambda kv: -kv[1]):
        if days >= d:
            return name
    return "1d"


def _next(seed: bytes) -> bytes:
    return hashlib.sha256(b"goldfish/next" + seed).digest()


@dataclass
class Keyring:
    se: SecureElement

    @staticmethod
    def slot(rung: int, cls: str) -> str:
        return f"R{rung}/{cls}"

    def provision(self, start_day: int) -> None:
        for r in RUNGS:
            for c in CLASSES:
                self.se.write_head(self.slot(r, c), start_day, os.urandom(32))

    def day_key(self, rung: int, cls: str, day: int) -> bytes:
        head_day, seed = self.se.read_head(self.slot(rung, cls))
        if day < head_day:
            raise KeyError(f"R{rung}/{cls} day {day} has been forgotten")
        for _ in range(day - head_day):
            seed = _next(seed)
        return self.se.bind(b"goldfish/day" + bytes([rung]) + cls.encode() + seed)

    def hour_key(self, rung: int, cls: str, day: int, hour: int) -> bytes:
        return hmac.new(self.day_key(rung, cls, day), b"hour" + bytes([hour]), hashlib.sha256).digest()

    def forget_through(self, rung: int, cls: str, day: int) -> None:
        """Advance the head so `day` and everything before it is unrecoverable."""
        s = self.slot(rung, cls)
        head_day, seed = self.se.read_head(s)
        while head_day <= day:
            seed = _next(seed)
            head_day += 1
        self.se.write_head(s, head_day, seed)

    def midnight(self, today: int, r3_days: int, r2_until_day: int, r1_days: int | None) -> None:
        """Apply the forgetting curve: one head write per chain at most."""
        for cls, cls_days in CLASSES.items():
            limits = {3: r3_days, 2: r2_until_day, 1: r1_days}
            for r in RUNGS:
                keep = limits[r]
                if cls_days is not None:
                    keep = cls_days if keep is None else min(keep, cls_days)
                if keep is not None:
                    self.forget_through(r, cls, today - keep)

    def fail_closed(self, today: int) -> None:
        """Clock can't be trusted: drop every R3 and R2 chain, keep outcomes."""
        for r in (2, 3):
            for c in CLASSES:
                self.forget_through(r, c, today + 10_000)