#![allow(unsafe_op_in_unsafe_fn)]

//! Trust store — verifies node identities against signed Merkle log checkpoints.
//!
//! [`TrustStore`] is the central piece of the NWEP trust model. It maintains a set
//! of trusted BLS anchor public keys and a collection of Merkle log checkpoints.
//! Given a [`MerkleEntry`] (an identity event) and its [`MerkleProof`], it can
//! verify that the entry is included in a checkpoint that has been signed by a
//! quorum of anchors — confirming that the identity is legitimately registered and
//! not revoked.
//!
//! Key concepts:
//! - **Checkpoint freshness**: [`check_staleness`](TrustStore::check_staleness) reports
//!   how old the latest checkpoint is. Clients should refuse to verify identities against
//!   checkpoints older than [`STALENESS_REJECT_NS`].
//! - **Equivocation detection**: [`check_equivocation`](TrustStore::check_equivocation)
//!   catches anchors that have signed conflicting checkpoints for the same epoch.
//! - **Identity caching**: [`cache_identity`](TrustStore::cache_identity) and
//!   [`lookup_identity`](TrustStore::lookup_identity) provide a short-term cache so
//!   repeated verifications for the same peer are fast.

use crate::ffi;
use crate::error::{check, Error};
use crate::types::{NodeId, Tstamp, Duration, MerkleHash};
use crate::bls::BlsPubkey;
use crate::anchor::AnchorSet;
use crate::checkpoint::Checkpoint;
use crate::merkle::{MerkleEntry, MerkleProof};

/// `STALENESS_WARNING_NS` is the checkpoint age (1 hour) at which [`Staleness::Warning`] is returned.
pub const STALENESS_WARNING_NS: Duration = 3600 * crate::types::SECONDS;
/// `STALENESS_REJECT_NS` is the checkpoint age (24 hours) at which [`Staleness::Reject`] is returned.
pub const STALENESS_REJECT_NS: Duration = 86400 * crate::types::SECONDS;
/// `IDENTITY_CACHE_TTL` is the default TTL (1 hour) for entries in the identity cache.
pub const IDENTITY_CACHE_TTL: Duration = 3600 * crate::types::SECONDS;

/// `Staleness` indicates how old the latest checkpoint in a [`TrustStore`] is.
///
/// Returned by [`TrustStore::check_staleness`]. Callers should at minimum warn on
/// `Warning` and refuse to accept new identities on `Reject`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Staleness {
    /// `Fresh` — the latest checkpoint is younger than [`STALENESS_WARNING_NS`].
    Fresh,
    /// `Warning` — the latest checkpoint is older than [`STALENESS_WARNING_NS`] but younger than [`STALENESS_REJECT_NS`].
    Warning,
    /// `Reject` — the latest checkpoint is older than [`STALENESS_REJECT_NS`]; identity verification should be refused.
    Reject,
}

impl From<ffi::nwep_staleness> for Staleness {
    fn from(s: ffi::nwep_staleness) -> Self {
        match s {
            ffi::nwep_staleness_NWEP_STALENESS_WARNING => Staleness::Warning,
            ffi::nwep_staleness_NWEP_STALENESS_REJECT => Staleness::Reject,
            _ => Staleness::Fresh,
        }
    }
}

/// `VerifiedIdentity` is the result of a successful [`TrustStore::verify_identity`] call.
///
/// It records the node's current public key, the log position and checkpoint epoch at
/// which it was verified, and whether it has since been revoked.
#[derive(Clone, Debug)]
pub struct VerifiedIdentity {
    /// The 32-byte node identifier.
    pub node_id: NodeId,
    /// The node's active Ed25519 public key at the time of verification.
    pub pubkey: [u8; 32],
    /// The Merkle log index of the entry that was verified.
    pub log_index: u64,
    /// The epoch of the checkpoint used to anchor this verification.
    pub checkpoint_epoch: u64,
    /// Timestamp (nanoseconds since epoch) when this identity was verified.
    pub verified_at: Tstamp,
    /// `true` if the node's identity has been revoked.
    pub revoked: bool,
}

impl VerifiedIdentity {
    pub(crate) fn from_ffi(v: &ffi::nwep_verified_identity) -> Self {
        VerifiedIdentity {
            node_id: NodeId(v.nodeid.data),
            pubkey: v.pubkey,
            log_index: v.log_index,
            checkpoint_epoch: v.checkpoint_epoch,
            verified_at: v.verified_at,
            revoked: v.revoked != 0,
        }
    }
}

/// `Equivocation` describes an anchor that has signed two conflicting checkpoints.
///
/// Returned by [`TrustStore::check_equivocation`] when misbehaviour is detected.
/// An equivocating anchor has signed two different Merkle roots for the same `epoch`,
/// which is a protocol violation; the anchor's key should be removed from the trust set.
#[derive(Clone, Debug)]
pub struct Equivocation {
    /// The BLS public key of the equivocating anchor.
    pub anchor: BlsPubkey,
    /// The epoch for which two conflicting signatures were found.
    pub epoch: u64,
    /// The first Merkle root signed by the anchor for this epoch.
    pub root1: MerkleHash,
    /// The second, conflicting Merkle root signed by the anchor for this epoch.
    pub root2: MerkleHash,
}

impl Equivocation {
    fn from_ffi(e: &ffi::nwep_equivocation) -> Self {
        Equivocation {
            anchor: BlsPubkey(e.anchor.data),
            epoch: e.epoch,
            root1: MerkleHash(e.root1.data),
            root2: MerkleHash(e.root2.data),
        }
    }
}

/// `TrustSettings` configures staleness thresholds and identity cache behaviour for a [`TrustStore`].
pub struct TrustSettings {
    /// Checkpoint age (nanoseconds) at which [`Staleness::Warning`] is returned. Default: [`STALENESS_WARNING_NS`].
    pub staleness_warning_ns: Duration,
    /// Checkpoint age (nanoseconds) at which [`Staleness::Reject`] is returned. Default: [`STALENESS_REJECT_NS`].
    pub staleness_reject_ns: Duration,
    /// Time-to-live (nanoseconds) for cached identity entries. Default: [`IDENTITY_CACHE_TTL`].
    pub identity_cache_ttl: Duration,
    /// Minimum number of anchor signatures required to validate a checkpoint.
    pub anchor_threshold: usize,
}

impl Default for TrustSettings {
    fn default() -> Self {
        TrustSettings {
            staleness_warning_ns: STALENESS_WARNING_NS,
            staleness_reject_ns: STALENESS_REJECT_NS,
            identity_cache_ttl: IDENTITY_CACHE_TTL,
            anchor_threshold: crate::types::DEFAULT_ANCHOR_THRESHOLD,
        }
    }
}

/// `TrustStorage` provides persistence callbacks for anchor keys and checkpoints.
///
/// Pass to [`TrustStore::new_with_storage`] to persist the trust store across
/// restarts. All fields are optional; unset callbacks are silently skipped.
pub struct TrustStorage {
    /// Called on startup to load the saved anchor public keys.
    pub anchor_load: Option<Box<dyn Fn() -> Result<Vec<BlsPubkey>, Error> + Send>>,
    /// Called when the anchor set changes, to persist the updated set.
    pub anchor_save: Option<Box<dyn Fn(&[BlsPubkey]) -> Result<(), Error> + Send>>,
    /// Called on startup to load previously saved checkpoints.
    pub checkpoint_load: Option<Box<dyn Fn() -> Result<Vec<Checkpoint>, Error> + Send>>,
    /// Called when a new checkpoint is added, to persist it durably.
    pub checkpoint_save: Option<Box<dyn Fn(&Checkpoint) -> Result<(), Error> + Send>>,
}

impl Default for TrustStorage {
    fn default() -> Self {
        TrustStorage {
            anchor_load: None,
            anchor_save: None,
            checkpoint_load: None,
            checkpoint_save: None,
        }
    }
}

struct StorageCallbacks {
    storage: TrustStorage,
}

unsafe extern "C" fn anchor_load_cb(
    user_data: *mut std::ffi::c_void,
    anchors: *mut ffi::nwep_bls_pubkey,
    max_anchors: usize,
) -> std::ffi::c_int {
    let cb = &*(user_data as *const StorageCallbacks);
    if let Some(f) = &cb.storage.anchor_load {
        match f() {
            Ok(pks) => {
                let n = pks.len().min(max_anchors);
                for (i, pk) in pks[..n].iter().enumerate() {
                    *anchors.add(i) = ffi::nwep_bls_pubkey { data: pk.0 };
                }
                n as i32
            }
            Err(e) => e.code,
        }
    } else {
        0
    }
}

unsafe extern "C" fn anchor_save_cb(
    user_data: *mut std::ffi::c_void,
    anchors: *const ffi::nwep_bls_pubkey,
    count: usize,
) -> std::ffi::c_int {
    let cb = &*(user_data as *const StorageCallbacks);
    if let Some(f) = &cb.storage.anchor_save {
        let pks: Vec<BlsPubkey> = (0..count).map(|i| BlsPubkey((*anchors.add(i)).data)).collect();
        match f(&pks) {
            Ok(()) => 0,
            Err(e) => e.code,
        }
    } else {
        0
    }
}

unsafe extern "C" fn checkpoint_load_cb(
    user_data: *mut std::ffi::c_void,
    checkpoints: *mut ffi::nwep_checkpoint,
    max_checkpoints: usize,
) -> std::ffi::c_int {
    let cb = &*(user_data as *const StorageCallbacks);
    if let Some(f) = &cb.storage.checkpoint_load {
        match f() {
            Ok(cps) => {
                let n = cps.len().min(max_checkpoints);
                for (i, cp) in cps[..n].iter().enumerate() {
                    *checkpoints.add(i) = cp.to_ffi();
                }
                n as i32
            }
            Err(e) => e.code,
        }
    } else {
        0
    }
}

unsafe extern "C" fn checkpoint_save_cb(
    user_data: *mut std::ffi::c_void,
    cp: *const ffi::nwep_checkpoint,
) -> std::ffi::c_int {
    let cb = &*(user_data as *const StorageCallbacks);
    if let Some(f) = &cb.storage.checkpoint_save {
        let checkpoint = Checkpoint::from_ffi(&*cp);
        match f(&checkpoint) {
            Ok(()) => 0,
            Err(e) => e.code,
        }
    } else {
        0
    }
}

/// `TrustStore` verifies node identities against a quorum of BLS-signed Merkle log checkpoints.
///
/// `TrustStore` holds a set of trusted anchor public keys and a collection of checkpoints.
/// Call [`verify_identity`](TrustStore::verify_identity) to confirm that a [`MerkleEntry`]
/// is covered by a valid checkpoint signed by enough anchors. Use
/// [`cache_identity`](TrustStore::cache_identity) to store the result so that subsequent
/// connections from the same peer are fast.
pub struct TrustStore {
    ptr: *mut ffi::nwep_trust_store,
    _callbacks: Option<Box<StorageCallbacks>>,
}

unsafe impl Send for TrustStore {}

impl TrustStore {
    /// `new` creates a `TrustStore` with in-memory state that is not persisted across restarts.
    ///
    /// Use [`new_with_storage`](TrustStore::new_with_storage) if persistence is required.
    ///
    /// # Errors
    ///
    /// Returns `Err` if the underlying C allocation fails.
    pub fn new(settings: TrustSettings) -> Result<Self, Error> {
        let ffi_settings = ffi::nwep_trust_settings {
            staleness_warning_ns: settings.staleness_warning_ns,
            staleness_reject_ns: settings.staleness_reject_ns,
            identity_cache_ttl: settings.identity_cache_ttl,
            anchor_threshold: settings.anchor_threshold,
        };
        let mut ptr: *mut ffi::nwep_trust_store = std::ptr::null_mut();
        check(unsafe { ffi::nwep_trust_store_new(&mut ptr, &ffi_settings, std::ptr::null()) })?;
        Ok(TrustStore { ptr, _callbacks: None })
    }

    /// `new_with_storage` creates a `TrustStore` backed by caller-provided persistence callbacks.
    ///
    /// On creation the C library calls `anchor_load` and `checkpoint_load` to restore
    /// previously saved state. Subsequent changes call `anchor_save` and `checkpoint_save`.
    ///
    /// # Errors
    ///
    /// Returns `Err` if the underlying C allocation fails.
    pub fn new_with_storage(settings: TrustSettings, storage: TrustStorage) -> Result<Self, Error> {
        let ffi_settings = ffi::nwep_trust_settings {
            staleness_warning_ns: settings.staleness_warning_ns,
            staleness_reject_ns: settings.staleness_reject_ns,
            identity_cache_ttl: settings.identity_cache_ttl,
            anchor_threshold: settings.anchor_threshold,
        };
        let mut cb = Box::new(StorageCallbacks { storage });
        let cb_ptr = cb.as_mut() as *mut _ as *mut std::ffi::c_void;
        let ffi_storage = ffi::nwep_trust_storage {
            anchor_load: Some(anchor_load_cb),
            anchor_save: Some(anchor_save_cb),
            checkpoint_load: Some(checkpoint_load_cb),
            checkpoint_save: Some(checkpoint_save_cb),
            user_data: cb_ptr,
        };
        let mut ptr: *mut ffi::nwep_trust_store = std::ptr::null_mut();
        check(unsafe { ffi::nwep_trust_store_new(&mut ptr, &ffi_settings, &ffi_storage) })?;
        Ok(TrustStore { ptr, _callbacks: Some(cb) })
    }

    /// `add_anchor` adds a BLS anchor public key to the trust store's anchor set.
    ///
    /// # Errors
    ///
    /// Returns `Err` if the key is already present or the C call fails.
    pub fn add_anchor(&mut self, pk: &BlsPubkey, builtin: bool) -> Result<(), Error> {
        let ffi_pk = ffi::nwep_bls_pubkey { data: pk.0 };
        check(unsafe { ffi::nwep_trust_store_add_anchor(self.ptr, &ffi_pk, builtin as i32) })
    }

    /// `remove_anchor` removes a BLS anchor public key from the trust store.
    ///
    /// # Errors
    ///
    /// Returns `Err` if the key is not present or the C call fails.
    pub fn remove_anchor(&mut self, pk: &BlsPubkey) -> Result<(), Error> {
        let ffi_pk = ffi::nwep_bls_pubkey { data: pk.0 };
        check(unsafe { ffi::nwep_trust_store_remove_anchor(self.ptr, &ffi_pk) })
    }

    /// `anchors` returns a shared reference to the internal [`AnchorSet`].
    ///
    /// The returned reference is valid for the lifetime of `self`.
    pub fn anchors(&self) -> &AnchorSet {
        unsafe {
            let ptr = ffi::nwep_trust_store_get_anchors(self.ptr);
            // SAFETY: The returned pointer is owned by the trust store and valid as long as self is.
            // We transmute the const pointer to a reference to AnchorSet which has the same layout.
            &*(ptr as *const AnchorSet)
        }
    }

    /// `add_checkpoint` adds a signed checkpoint to the trust store.
    ///
    /// The checkpoint is validated against the current anchor set before being stored.
    ///
    /// # Errors
    ///
    /// Returns `Err` if the checkpoint signature is invalid or the C call fails.
    pub fn add_checkpoint(&mut self, cp: &Checkpoint) -> Result<(), Error> {
        let ffi_cp = cp.to_ffi();
        check(unsafe { ffi::nwep_trust_store_add_checkpoint(self.ptr, &ffi_cp) })
    }

    /// `get_latest_checkpoint` returns the most recently stored checkpoint.
    ///
    /// # Errors
    ///
    /// Returns `Err` if no checkpoints have been stored yet.
    pub fn get_latest_checkpoint(&self) -> Result<Checkpoint, Error> {
        let mut cp = unsafe { std::mem::zeroed::<ffi::nwep_checkpoint>() };
        check(unsafe { ffi::nwep_trust_store_get_latest_checkpoint(self.ptr, &mut cp) })?;
        Ok(Checkpoint::from_ffi(&cp))
    }

    /// `get_checkpoint` retrieves the checkpoint for a specific epoch.
    ///
    /// # Errors
    ///
    /// Returns `Err` if no checkpoint for `epoch` has been stored.
    pub fn get_checkpoint(&self, epoch: u64) -> Result<Checkpoint, Error> {
        let mut cp = unsafe { std::mem::zeroed::<ffi::nwep_checkpoint>() };
        check(unsafe { ffi::nwep_trust_store_get_checkpoint(self.ptr, epoch, &mut cp) })?;
        Ok(Checkpoint::from_ffi(&cp))
    }

    /// `checkpoint_count` returns the number of checkpoints stored in the trust store.
    pub fn checkpoint_count(&self) -> usize {
        unsafe { ffi::nwep_trust_store_checkpoint_count(self.ptr) }
    }

    /// `check_staleness` reports how old the latest stored checkpoint is relative to `now`.
    ///
    /// Returns [`Staleness::Fresh`] if checkpoints are recent, [`Staleness::Warning`] if
    /// older than [`STALENESS_WARNING_NS`], or [`Staleness::Reject`] if older than
    /// [`STALENESS_REJECT_NS`].
    pub fn check_staleness(&self, now: Tstamp) -> Staleness {
        Staleness::from(unsafe { ffi::nwep_trust_store_check_staleness(self.ptr, now) })
    }

    /// `get_staleness_age` returns the age (nanoseconds) of the latest checkpoint relative to `now`.
    pub fn get_staleness_age(&self, now: Tstamp) -> Duration {
        unsafe { ffi::nwep_trust_store_get_staleness_age(self.ptr, now) }
    }

    /// `verify_identity` verifies that a Merkle log entry is covered by a trusted, quorum-signed checkpoint.
    ///
    /// Checks that `proof` is a valid inclusion proof for `entry` in a checkpoint with a root
    /// known to the trust store, and that the checkpoint has sufficient anchor signatures.
    /// If `checkpoint` is `Some`, that specific checkpoint is used; otherwise the best known
    /// checkpoint is selected automatically.
    ///
    /// # Errors
    ///
    /// Returns `Err` if the proof is invalid, no suitable checkpoint exists, the checkpoint
    /// lacks sufficient signatures, or the identity has been revoked.
    pub fn verify_identity(
        &mut self,
        entry: &MerkleEntry,
        proof: &MerkleProof,
        checkpoint: Option<&Checkpoint>,
        now: Tstamp,
    ) -> Result<VerifiedIdentity, Error> {
        let ffi_entry = entry.to_ffi();
        let ffi_proof = proof.to_ffi();
        let ffi_cp_owned;
        let ffi_cp_ptr = match checkpoint {
            Some(cp) => { ffi_cp_owned = cp.to_ffi(); &ffi_cp_owned as *const _ }
            None => std::ptr::null(),
        };
        let mut result = unsafe { std::mem::zeroed::<ffi::nwep_verified_identity>() };
        check(unsafe {
            ffi::nwep_trust_store_verify_identity(self.ptr, &ffi_entry, &ffi_proof, ffi_cp_ptr, now, &mut result)
        })?;
        Ok(VerifiedIdentity::from_ffi(&result))
    }

    /// `cache_identity` stores a previously verified identity for fast future lookups.
    ///
    /// The cached entry expires after `identity_cache_ttl` nanoseconds.
    ///
    /// # Errors
    ///
    /// Returns `Err` if the C call fails.
    pub fn cache_identity(&mut self, vi: &VerifiedIdentity) -> Result<(), Error> {
        let ffi_vi = ffi::nwep_verified_identity {
            nodeid: ffi::nwep_nodeid { data: vi.node_id.0 },
            pubkey: vi.pubkey,
            log_index: vi.log_index,
            checkpoint_epoch: vi.checkpoint_epoch,
            verified_at: vi.verified_at,
            revoked: vi.revoked as i32,
        };
        check(unsafe { ffi::nwep_trust_store_cache_identity(self.ptr, &ffi_vi) })
    }

    /// `lookup_identity` retrieves a cached verified identity if it has not expired.
    ///
    /// # Errors
    ///
    /// Returns `Err` if no valid cached entry exists for `node_id`.
    pub fn lookup_identity(&self, node_id: &NodeId, now: Tstamp) -> Result<VerifiedIdentity, Error> {
        let ffi_nid = ffi::nwep_nodeid { data: node_id.0 };
        let mut out = unsafe { std::mem::zeroed::<ffi::nwep_verified_identity>() };
        check(unsafe { ffi::nwep_trust_store_lookup_identity(self.ptr, &ffi_nid, now, &mut out) })?;
        Ok(VerifiedIdentity::from_ffi(&out))
    }

    /// `check_equivocation` detects whether an anchor has signed two conflicting checkpoints for the same epoch.
    ///
    /// Returns `Ok(Some(equivocation))` if misbehaviour is detected, `Ok(None)` if the checkpoint
    /// is consistent with previously seen checkpoints, or `Err` on a C-level failure.
    pub fn check_equivocation(&mut self, cp: &Checkpoint) -> Result<Option<Equivocation>, Error> {
        let ffi_cp = cp.to_ffi();
        let mut out = unsafe { std::mem::zeroed::<ffi::nwep_equivocation>() };
        let rc = unsafe { ffi::nwep_trust_store_check_equivocation(self.ptr, &ffi_cp, &mut out) };
        if rc == 0 {
            Ok(None)
        } else if rc == crate::error::ERR_TRUST_EQUIVOCATION {
            Ok(Some(Equivocation::from_ffi(&out)))
        } else {
            Err(Error::from_code(rc))
        }
    }
}

impl Drop for TrustStore {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            unsafe { ffi::nwep_trust_store_free(self.ptr) }
        }
    }
}
