//! The cube environment: how Null cubes select key material + transform, and
//! how records are sealed/opened against that environment (PDF Package 5).
//!
//! The PDF's core idea: "exploit Null cubes/layers as environment settings
//! that influence how records are read/written." Concretely:
//!
//! * A [`CubeEnv`] holds one or more [`KeySlot`]s. Each slot points at a
//!   Null-cube cell whose *body* is raw key material, and says which
//!   [`TransformId`] to apply.
//! * A record is encrypted under a slot selected by a *selector* (a tenant
//!   byte, or derived from the record's own coordinate). Because the key
//!   material physically lives in a cube cell, "same CZYX coordinate means
//!   different plaintext under different Null settings" is literal: re-point
//!   the env's slot at a different Null cube and the same coordinate opens to
//!   different data.
//! * Sealing/opening stamp and read a [`HEADER_FLAG_ENCRYPTED`] bit on the
//!   record's [`CubeHeader`], so the store can tell an encrypted record from
//!   a plaintext one without guessing.

use cubecoords::{CubeHeader, Czyx};
use cubestore::{CubeBackend, CubeStore};

use crate::transform::{self, CryptoError, Key, KeySlot};

/// Header flag bit (out-of-band, stored in `CubeHeader.flags` spare range)
/// marking a record body as a cubecrypt envelope.
pub const HEADER_FLAG_ENCRYPTED: u16 = 1 << 12;

/// Selects which [`KeySlot`]/transform applies to a record.
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum Selector {
    /// Explicit tenant/slot index into the env's slot list.
    Slot(u8),
    /// Derive the slot from a record coordinate (uses `c` axis as the slot
    /// index, mod slot count). Lets the same env key records per-class.
    FromCoord(Czyx),
}

/// A cube crypto environment: key material in Null cubes + the transforms they
/// select.
pub struct CubeEnv {
    slots: Vec<KeySlot>,
    /// Salt mixed into every key derivation (env-wide).
    env_salt: Vec<u8>,
}

impl CubeEnv {
    /// Build an environment from a list of key slots.
    pub fn new(slots: Vec<KeySlot>, env_salt: Vec<u8>) -> Self {
        CubeEnv { slots, env_salt }
    }

    /// Number of key slots.
    pub fn len(&self) -> usize {
        self.slots.len()
    }

    /// True when there are no slots (a plaintext-only env).
    pub fn is_empty(&self) -> bool {
        self.slots.is_empty()
    }

    /// Resolve which slot a selector maps to.
    fn slot_for(&self, sel: Selector) -> Option<&KeySlot> {
        match sel {
            Selector::Slot(i) => self.slots.get(i as usize),
            Selector::FromCoord(c) => {
                if self.slots.is_empty() {
                    None
                } else {
                    let idx = (c.c as usize) % self.slots.len();
                    Some(&self.slots[idx])
                }
            }
        }
    }

    /// Resolve the concrete key bytes for a slot by reading its Null-cube cell
    /// from the store.
    fn key_for<B: CubeBackend>(
        &self,
        store: &CubeStore<B>,
        slot: &KeySlot,
    ) -> Result<Key, EnvError> {
        let (_, body) = store
            .get_record(&slot.key_cell)
            .ok_or(EnvError::KeyCellMissing(slot.key_cell))?;
        if body.len() < 16 {
            return Err(EnvError::KeyMaterialTooShort(slot.key_cell));
        }
        let mut salt = self.env_salt.clone();
        salt.extend_from_slice(&slot.salt);
        Ok(transform::derive_key(&body, &salt))
    }

    /// Encrypt `plaintext` for a record, selected by `sel`, returning a
    /// cubecrypt envelope. The returned header has the encrypted bit set and
    /// `doc_type` optionally re-tagged so tools can spot ciphertext.
    pub fn seal<B: CubeBackend>(
        &self,
        store: &CubeStore<B>,
        sel: Selector,
        plaintext: &[u8],
    ) -> Result<Vec<u8>, EnvError> {
        let slot = self.slot_for(sel).ok_or(EnvError::NoSlotFor(sel))?;
        let key = self.key_for(store, slot)?;
        Ok(transform::seal(slot.transform, &key, plaintext))
    }

    /// Decrypt a cubecrypt envelope previously produced by [`seal`].
    pub fn open<B: CubeBackend>(
        &self,
        store: &CubeStore<B>,
        sel: Selector,
        envelope: &[u8],
    ) -> Result<Vec<u8>, EnvError> {
        let slot = self.slot_for(sel).ok_or(EnvError::NoSlotFor(sel))?;
        let key = self.key_for(store, slot)?;
        transform::open(&key, envelope).map_err(EnvError::Crypto)
    }

    /// Encrypt `plaintext` and write it as a record at `label`, returning the
    /// header to store alongside it (encrypted flag set).
    pub fn put_encrypted<B: CubeBackend>(
        &self,
        store: &mut CubeStore<B>,
        label: Czyx,
        sel: Selector,
        plaintext: &[u8],
        mut header: CubeHeader,
    ) -> Result<(), EnvError> {
        let envelope = self.seal(store, sel, plaintext)?;
        header.flags.0 |= HEADER_FLAG_ENCRYPTED;
        if header.doc_type.is_none() {
            header.doc_type = Some("cubecrypt".into());
        }
        header.size_bytes = Some(plaintext.len() as u64);
        header.refresh_flags();
        store.put_record(label, &header, &envelope);
        Ok(())
    }

    /// Read an encrypted record at `label` and return its plaintext.
    pub fn get_decrypted<B: CubeBackend>(
        &self,
        store: &CubeStore<B>,
        label: Czyx,
        sel: Selector,
    ) -> Result<Vec<u8>, EnvError> {
        let (_, envelope) = store
            .get_record(&label)
            .ok_or(EnvError::RecordMissing(label))?;
        self.open(store, sel, &envelope)
    }
}

/// Errors from environment resolution / record crypto.
#[derive(Clone, Eq, PartialEq, Debug)]
pub enum EnvError {
    /// No key slot matched the selector.
    NoSlotFor(Selector),
    /// The Null-cube cell holding key material is absent.
    KeyCellMissing(Czyx),
    /// Key material too short to derive a key.
    KeyMaterialTooShort(Czyx),
    /// The target record does not exist.
    RecordMissing(Czyx),
    /// Underlying crypto failure (bad envelope / auth fail / feature missing).
    Crypto(CryptoError),
}