"""Persist and reload the clean-room crypto authority across demo commands.

The walkthrough runs package (05), serve (06) and revoke (07) as separate
processes. The certificate authority, revocation ledger, wrapping-key custody,
sealed unit and the licensee's leaf certificate/key are therefore written to a
gitignored ``state/`` directory between steps.

Everything here uses only the repository's own ``capability_licensing`` API plus
its public ``export_state`` / ``load_state`` methods. The wrapping keys in
``keystore.json`` are demo-grade custody, which is exactly why ``state/`` is
gitignored and never committed.
"""

from __future__ import annotations

import json
from dataclasses import asdict
from pathlib import Path
from typing import Tuple

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

# Import .common first: it puts src/ on sys.path so capability_licensing resolves.
from .common import STATE_DIR

from capability_licensing import (
    Certificate,
    CertificateAuthority,
    KeyAuthority,
    RevocationList,
    SealedUnit,
)

_CA = STATE_DIR / "ca.json"
_REVOCATION = STATE_DIR / "revocation.json"
_KEYSTORE = STATE_DIR / "keystore.json"
_UNIT = STATE_DIR / "unit.json"
_EXFIL = STATE_DIR / "exfiltrated_unit.json"
_LEAF_CERT = STATE_DIR / "leaf_cert.json"
_LEAF_KEY = STATE_DIR / "leaf_key.hex"
_META = STATE_DIR / "meta.json"


def _write_json(path: Path, data) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2))


def _read_json(path: Path):
    return json.loads(path.read_text())


# -- certificate authority + revocation -----------------------------------


def save_ca(ca: CertificateAuthority) -> None:
    _write_json(_CA, ca.export_state())


def load_ca() -> CertificateAuthority:
    return CertificateAuthority.restore(_read_json(_CA))


def save_revocation(revocation: RevocationList) -> None:
    _write_json(_REVOCATION, revocation.export_state())


def load_revocation() -> RevocationList:
    return RevocationList.restore(_read_json(_REVOCATION))


# -- wrapping-key custody --------------------------------------------------


def save_keystore(keys: KeyAuthority) -> None:
    _write_json(_KEYSTORE, keys.export_state())


def load_keystore(ca: CertificateAuthority, revocation: RevocationList) -> KeyAuthority:
    meta = _read_json(_META)
    root = ca.certificate(meta["root_serial"])
    organisation = ca.certificate(meta["org_serial"])
    keys = KeyAuthority(
        trusted_roots=[root],
        intermediates=[organisation],
        revocation_list=revocation,
    )
    keys.load_state(_read_json(_KEYSTORE))
    return keys


# -- sealed units ----------------------------------------------------------


def save_unit(unit: SealedUnit) -> None:
    _write_json(_UNIT, asdict(unit))


def load_unit() -> SealedUnit:
    return SealedUnit(**_read_json(_UNIT))


def save_exfiltrated_copy(unit: SealedUnit) -> None:
    """An attacker's byte-for-byte copy of the encrypted unit, held elsewhere."""
    _write_json(_EXFIL, asdict(unit))


def load_exfiltrated_copy() -> SealedUnit:
    return SealedUnit(**_read_json(_EXFIL))


# -- leaf licence ----------------------------------------------------------


def save_leaf(certificate: Certificate, private_key: Ed25519PrivateKey) -> None:
    _write_json(_LEAF_CERT, certificate.to_dict())
    _LEAF_KEY.write_text(private_key.private_bytes_raw().hex())


def load_leaf() -> Tuple[Certificate, Ed25519PrivateKey]:
    certificate = Certificate(**_read_json(_LEAF_CERT))
    private_key = Ed25519PrivateKey.from_private_bytes(bytes.fromhex(_LEAF_KEY.read_text()))
    return certificate, private_key


def save_meta(root_serial: str, org_serial: str, leaf_serial: str, unit_id: str) -> None:
    _write_json(
        _META,
        {
            "root_serial": root_serial,
            "org_serial": org_serial,
            "leaf_serial": leaf_serial,
            "unit_id": unit_id,
        },
    )


def load_meta():
    return _read_json(_META)