"""HKDF key hierarchy for crypto-shred erasure (A4/A8). k_org_master → k_subject → k_record.

Ported from the NCTE adapter (Compliance-Safe §2) into the Ledger core's A8 encrypted-external slot.

CRYPTO-SHREDDING ERASURE. Derived keys are EPHEMERAL: `k_subject` and `k_record` are derived on demand and
live only for one operation; they are NEVER persisted or cached. The only durable per-subject secret is a
random `subject_salt` in the `SubjectKeyStore`; `k_subject` is HKDF-derived from `k_org_master` WITH that
salt, so destroying the salt makes `k_subject` (and every `k_record` and opening beneath it) un-derivable
even with the master — the cryptographic basis for subject-level erasure. After `destroy_subject(uid)` no
code path can produce an opening for that subject's records.

Discipline:
  * `subject_uid` / `record_uid` are OPAQUE internal ids — never an email/name (PII in a derivation `info`
    string would couple key management to personal data; the uid↔person map is a SEPARATE record).
  * The org master is reached through a `KeyManager` INTERFACE — a real KMS/HSM is deployment config, not
    code. Destruction goes through a `DualControl` hook and is recorded in an append-only lifecycle audit.
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Protocol

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

KEY_LEN = 32
SUBJECT_INFO_PREFIX = b"subject:"
RECORD_INFO_PREFIX = b"record:"


class ErasedError(Exception):
    """Raised when an opening is requested for a subject whose key material has been destroyed."""


class Erasability:
    SUBJECT_LEVEL = "subject_level"      # minted under the hierarchy — destroying k_subject erases all records
    PER_RECORD_ONLY = "per_record_only"  # legacy flat-key — no subject-level erasure can be promised


def subject_erasability(erasabilities: List[str]) -> str:
    """Weakest-link: a subject is subject-erasable only if EVERY backing record is. One legacy
    (`per_record_only`) record caps the subject to `per_record_only`."""
    return (
        Erasability.SUBJECT_LEVEL
        if erasabilities and all(e == Erasability.SUBJECT_LEVEL for e in erasabilities)
        else Erasability.PER_RECORD_ONLY
    )


class KeyManager(Protocol):
    """Interface to the org master key. A real KMS/HSM is deployment config; tests inject an in-memory one."""

    def org_master(self) -> bytes: ...


@dataclass
class InMemoryKeyManager:
    """Test/dev KeyManager. In production this is a KMS handle, not raw bytes."""

    _master: bytes

    def org_master(self) -> bytes:
        return self._master


DualControl = Callable[[str, str], bool]  # (action, subject_uid) -> approved


def _single_actor_denied(action: str, subject_uid: str) -> bool:
    """Default dual-control hook: deny. Destruction MUST be explicitly dual-authorised."""
    return False


@dataclass
class KeyLifecycleAudit:
    """Append-only audit of key lifecycle events (generate / derive / destroy)."""

    events: List[dict] = field(default_factory=list)

    def record(self, event: str, *, subject_uid: Optional[str] = None, trigger_ref: Optional[str] = None) -> None:
        self.events.append({"event": event, "subject_uid": subject_uid, "trigger_ref": trigger_ref})


@dataclass
class SubjectKeyStore:
    """Holds the random per-subject salts. Destroying a subject's salt crypto-shreds the subject:
    `k_subject` can no longer be derived even with the master."""

    _salts: Dict[str, bytes] = field(default_factory=dict)

    def ensure(self, uid: str, *, salt: Optional[bytes] = None) -> bytes:
        if uid not in self._salts:
            self._salts[uid] = salt if salt is not None else os.urandom(KEY_LEN)
        return self._salts[uid]

    def get(self, uid: str) -> Optional[bytes]:
        return self._salts.get(uid)

    def exists(self, uid: str) -> bool:
        return uid in self._salts

    def destroy(self, uid: str) -> None:
        self._salts.pop(uid, None)


def _hkdf(ikm: bytes, *, salt: Optional[bytes], info: bytes) -> bytes:
    return HKDF(algorithm=hashes.SHA256(), length=KEY_LEN, salt=salt, info=info).derive(ikm)


class KeyHierarchy:
    """Derives k_subject / k_record on demand (ephemeral) and crypto-shreds on destroy. Holds NO derived
    keys — only the KMS handle, the salt store, the audit log, and the dual-control hook."""

    def __init__(
        self,
        key_manager: KeyManager,
        store: SubjectKeyStore,
        audit: KeyLifecycleAudit,
        *,
        dual_control: DualControl = _single_actor_denied,
    ) -> None:
        self._km = key_manager
        self._store = store
        self._audit = audit
        self._dual_control = dual_control

    def generate_subject(self, uid: str, *, salt: Optional[bytes] = None) -> None:
        """Create the destroyable per-subject secret (returns no key)."""
        self._store.ensure(uid, salt=salt)
        self._audit.record("generate", subject_uid=uid)

    def derive_subject_key(self, uid: str) -> bytes:
        """Ephemeral k_subject = HKDF(ikm=k_org_master, salt=subject_salt, info='subject:'||uid). Raises
        ErasedError if the subject's salt has been destroyed."""
        salt = self._store.get(uid)
        if salt is None:
            raise ErasedError(f"subject {uid!r} is erased — k_subject is not derivable")
        self._audit.record("derive", subject_uid=uid)
        return _hkdf(self._km.org_master(), salt=salt, info=SUBJECT_INFO_PREFIX + uid.encode("utf-8"))

    def derive_record_key(self, uid: str, record_uid: str) -> bytes:
        """Ephemeral k_record = HKDF(ikm=k_subject, salt=None, info='record:'||record_uid)."""
        k_subject = self.derive_subject_key(uid)  # raises if erased
        return _hkdf(k_subject, salt=None, info=RECORD_INFO_PREFIX + record_uid.encode("utf-8"))

    def open_record(self, uid: str, record_uid: str) -> bytes:
        """Produce the opening (k_record) for a record. After the subject is destroyed this raises
        ErasedError — there is no other code path to an opening."""
        return self.derive_record_key(uid, record_uid)

    def destroy_subject(self, uid: str, *, trigger_ref: Optional[str] = None) -> None:
        """Crypto-shred a subject: destroy its salt (dual-control required) and audit it. After this, no
        opening for the subject's records can be derived."""
        if not self._dual_control("destroy_subject", uid):
            raise PermissionError("destroy_subject requires dual-control authorisation")
        self._store.destroy(uid)
        self._audit.record("destroy", subject_uid=uid, trigger_ref=trigger_ref)
