//! Checkpoint and restore for the `SignalLedger`.
//!
//! # Checkpoint
//!
//! `SignalLedger::checkpoint()` serializes all in-memory signal state to the
//! `StorageEngine` as a single atomic `WriteBatch`. No partial checkpoints are
//! possible: either the whole ledger is written or nothing is.
//!
//! # Restore
//!
//! `SignalLedger::restore()` scans the storage, filters for `Tag::Sig` keys,
//! deserializes each entry, and populates the `DashMap`. Returns the checkpoint
//! metadata (for WAL replay) or `None` if no checkpoint exists (first boot).
//!
//! # Binary format
//!
//! Each entry serializes as a 983-byte fixed-length record (see [`format`]).
//! The checkpoint metadata serializes as a 49-byte V2 record at a well-known key
//! (see [`meta`]). All payload values use little-endian byte order; storage
//! keys use big-endian (the existing `encode_key` convention). A version byte
//! at offset 0 enables future backward-compatible format changes.
//!
//! # Integrity
//!
//! V2 checkpoints include a BLAKE3 hash of all serialized entry payloads.
//! On restore, the hash is verified; a mismatch returns `Ok(None)` which
//! triggers a full WAL replay (fallback to safe state).

pub mod format;
pub(super) mod integrity;
pub mod meta;

pub use meta::CheckpointMeta;

use crate::schema::{EntityId, TidalError};
use crate::storage::{StorageEngine, Tag, WriteBatch, encode_key, parse_key};

use format::{deserialize_entry, serialize_entry};
use integrity::{hash_checkpoint_payload, verify_checkpoint_payload};
use meta::{META_SUFFIX, deserialize_meta, serialize_meta};

use super::ledger::SignalLedger;

// ── SignalLedger impl ─────────────────────────────────────────────────────────

impl SignalLedger {
    /// Write all in-memory signal state to the storage engine atomically.
    ///
    /// Iterates the `DashMap` and serializes each entry into a `WriteBatch`.
    /// The checkpoint metadata is stored at a well-known key:
    /// `encode_key(EntityId::new(0), Tag::Sig, b"meta")`.
    ///
    /// The BLAKE3 hash of all serialized entry values (collected in iteration
    /// order) is stored in `meta.payload_hash`. On restore, the hash is
    /// verified to detect silent data corruption.
    ///
    /// # Errors
    ///
    /// Returns `TidalError::Storage` if the `WriteBatch` commit or `flush` fails.
    ///
    /// # Concurrency
    ///
    /// This method iterates `DashMap` shards without a global lock. Entries
    /// written concurrently to already-snapshotted shards will be absent from
    /// the checkpoint. The caller must supply `meta.wal_sequence` equal to the
    /// WAL tail at checkpoint start; restore must replay from that sequence to
    /// recover any missing entries.
    pub fn checkpoint(
        &self,
        storage: &dyn StorageEngine,
        meta: CheckpointMeta,
    ) -> crate::Result<()> {
        let mut batch = WriteBatch::new();

        // Collect all (key, value) pairs for deterministic ordering.
        // DashMap iteration order is non-deterministic (hash-based sharding),
        // but storage scan_prefix returns keys in lexicographic order. We sort
        // by key so that the hash computed here matches the scan order on restore.
        let mut kv_pairs: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        for entry_ref in self.entries() {
            let &(entity_id, signal_type_id) = entry_ref.key();
            let entry = entry_ref.value();
            // Entry key suffix is the signal_type_id as 2 big-endian bytes,
            // so it is exactly 2 bytes -- never collides with b"meta" (4 bytes).
            let suffix = signal_type_id.as_u16().to_be_bytes();
            let key = encode_key(entity_id, Tag::Sig, &suffix);
            let value = serialize_entry(entity_id, signal_type_id, entry);
            kv_pairs.push((key, value));
        }
        kv_pairs.sort_by(|(a, _), (b, _)| a.cmp(b));

        // Extract sorted values for hashing, then add all to the batch.
        let entry_values: Vec<Vec<u8>> = kv_pairs.iter().map(|(_, v)| v.clone()).collect();
        for (key, value) in kv_pairs {
            batch.put(key, value);
        }

        // Compute BLAKE3 hash over entry payloads in sorted key order.
        let payload_hash = hash_checkpoint_payload(&entry_values);

        // Write checkpoint metadata with the computed hash.
        let meta_with_hash = CheckpointMeta {
            checkpoint_time_ns: meta.checkpoint_time_ns,
            wal_sequence: meta.wal_sequence,
            payload_hash,
        };
        let meta_key = encode_key(EntityId::new(0), Tag::Sig, META_SUFFIX);
        batch.put(meta_key, serialize_meta(&meta_with_hash));

        #[cfg(any(test, feature = "test-utils"))]
        crate::testing::crash_injector::check_crash_point(
            crate::testing::CrashPoint::CheckpointPreFlush,
        );

        storage.write_batch(batch)?;
        storage.flush()?;

        #[cfg(any(test, feature = "test-utils"))]
        crate::testing::crash_injector::check_crash_point(
            crate::testing::CrashPoint::CheckpointPostFlush,
        );

        Ok(())
    }

    /// Restore in-memory signal state from the storage engine.
    ///
    /// Scans all keys, filters for `Tag::Sig` entries (excluding the meta key),
    /// collects the raw values, and verifies the BLAKE3 integrity hash before
    /// deserializing. If the hash does not match (indicating corruption), the
    /// method logs a warning and returns `Ok(None)`, which triggers a full WAL
    /// replay as a safe fallback.
    ///
    /// Returns `Some(CheckpointMeta)` if a checkpoint exists, or `None` on
    /// first boot (empty storage) or on integrity verification failure.
    ///
    /// # Errors
    ///
    /// - `TidalError::Storage` on I/O failure
    /// - `TidalError::Internal` on deserialization failure (corrupt checkpoint)
    pub fn restore(&self, storage: &dyn StorageEngine) -> crate::Result<Option<CheckpointMeta>> {
        // Read checkpoint metadata first.
        let meta_key = encode_key(EntityId::new(0), Tag::Sig, META_SUFFIX);
        let meta = match storage.get(&meta_key)? {
            None => None,
            Some(meta_bytes) => Some(deserialize_meta(&meta_bytes).map_err(|e| {
                TidalError::internal(
                    "checkpoint_restore",
                    format!("corrupt checkpoint meta: {e}"),
                )
            })?),
        };

        // Collect raw entry values for integrity verification BEFORE deserializing.
        // scan_prefix returns keys in lexicographic order (same order used by
        // checkpoint() after sorting), so the hash input order is deterministic.
        let mut raw_values: Vec<Vec<u8>> = Vec::new();

        // TECH DEBT: scan_prefix(&[]) iterates the entire keyspace. This is safe
        // today (signals are the only key namespace), but must be replaced with a
        // Tag::Sig-scoped scan (e.g. `scan_tag(Tag::Sig)`) once M1P5 adds entity,
        // index, and embedding key namespaces to avoid iterating unrelated data.
        for item in storage.scan_prefix(&[]) {
            let (key, value) = item?;
            if let Some((entity_id, Tag::Sig, suffix)) = parse_key(&key) {
                // Skip the meta key (entity_id=0, suffix=b"meta").
                if entity_id == EntityId::new(0) && suffix == META_SUFFIX {
                    continue;
                }
                raw_values.push(value);
            }
        }

        // Verify integrity before deserializing or populating the DashMap.
        if let Some(ref m) = meta
            && !verify_checkpoint_payload(&raw_values, &m.payload_hash)
        {
            tracing::warn!(
                "checkpoint BLAKE3 integrity check failed; \
                 falling back to full WAL replay"
            );
            return Ok(None);
        }

        // Integrity verified (or V1 with zero hash); deserialize and insert.
        for value in &raw_values {
            let (eid, stid, entry) = deserialize_entry(value).map_err(|e| {
                TidalError::internal(
                    "checkpoint_restore",
                    format!("corrupt checkpoint entry: {e}"),
                )
            })?;
            self.entries.insert((eid, stid), entry);
        }

        Ok(meta)
    }

    /// Return the number of entries currently in the `DashMap`.
    ///
    /// Used for diagnostics and testing.
    #[must_use]
    pub fn entry_count(&self) -> usize {
        self.entries.len()
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use std::time::Duration;

    use super::*;
    use crate::schema::{DecaySpec, EntityKind, SchemaBuilder, Timestamp, Window};
    use crate::signals::ledger::NoopWalWriter;
    use crate::storage::InMemoryBackend;

    fn test_schema() -> crate::schema::Schema {
        let mut builder = SchemaBuilder::new();
        let _ = builder
            .signal(
                "view",
                EntityKind::Item,
                DecaySpec::Exponential {
                    half_life: Duration::from_secs(7 * 24 * 3600),
                },
            )
            .windows(&[Window::OneHour, Window::AllTime])
            .velocity(true)
            .add();
        builder.build().expect("valid test schema")
    }

    #[test]
    fn checkpoint_to_empty_storage() {
        let schema = test_schema();
        let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));

        let now = Timestamp::now();
        for i in 0..10u64 {
            ledger
                .record_signal("view", EntityId::new(i + 1), 1.0, now)
                .expect("record");
        }

        let storage = InMemoryBackend::new();
        let meta = CheckpointMeta {
            checkpoint_time_ns: now.as_nanos(),
            wal_sequence: 100,
            payload_hash: [0u8; 32], // hash computed by checkpoint()
        };
        ledger.checkpoint(&storage, meta).expect("checkpoint");

        // Expect meta key + 10 entry keys = 11 total keys.
        let all_keys: Vec<_> = storage
            .scan_prefix(&[])
            .collect::<Result<_, _>>()
            .expect("scan ok");
        assert_eq!(
            all_keys.len(),
            11,
            "expected 11 keys, got {}",
            all_keys.len()
        );
    }

    #[test]
    fn restore_from_empty_storage() {
        let schema = test_schema();
        let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));

        let storage = InMemoryBackend::new();
        let meta = ledger.restore(&storage).expect("restore ok");

        assert!(meta.is_none(), "empty storage should return None");
        assert_eq!(ledger.entry_count(), 0);
    }

    #[test]
    fn restore_preserves_decay_scores() {
        let schema = test_schema();
        let ledger = SignalLedger::new(schema.clone(), Box::new(NoopWalWriter));

        let ts1 = Timestamp::from_nanos(1_000_000_000_000);
        let ts2 = Timestamp::from_nanos(1_001_000_000_000);
        ledger
            .record_signal("view", EntityId::new(42), 5.0, ts1)
            .expect("record 1");
        ledger
            .record_signal("view", EntityId::new(42), 3.0, ts2)
            .expect("record 2");

        let storage = InMemoryBackend::new();
        let meta = CheckpointMeta {
            checkpoint_time_ns: 1_002_000_000_000,
            wal_sequence: 50,
            payload_hash: [0u8; 32], // hash computed by checkpoint()
        };
        ledger.checkpoint(&storage, meta).expect("checkpoint");

        let ledger2 = SignalLedger::new(schema, Box::new(NoopWalWriter));
        let restored_meta = ledger2
            .restore(&storage)
            .expect("restore")
            .expect("some meta");
        assert_eq!(restored_meta.wal_sequence, 50);

        let score_orig = ledger
            .read_decay_score(EntityId::new(42), "view", 0)
            .expect("ok");
        let score_rest = ledger2
            .read_decay_score(EntityId::new(42), "view", 0)
            .expect("ok");

        assert!(score_orig.is_some());
        assert!(score_rest.is_some());
    }

    #[test]
    fn restore_preserves_windowed_counts() {
        let schema = test_schema();
        let ledger = SignalLedger::new(schema.clone(), Box::new(NoopWalWriter));

        let base_ns = 1_000_000_000_000u64;
        for i in 0..100u64 {
            let ts = Timestamp::from_nanos(base_ns + i * 100_000_000);
            ledger
                .record_signal("view", EntityId::new(1), 1.0, ts)
                .expect("record");
        }

        let storage = InMemoryBackend::new();
        let meta = CheckpointMeta {
            checkpoint_time_ns: base_ns + 10_000_000_000,
            wal_sequence: 0,
            payload_hash: [0u8; 32], // hash computed by checkpoint()
        };
        ledger.checkpoint(&storage, meta).expect("checkpoint");

        let ledger2 = SignalLedger::new(schema, Box::new(NoopWalWriter));
        ledger2.restore(&storage).expect("restore");

        let count_orig = ledger
            .read_windowed_count(EntityId::new(1), "view", Window::AllTime)
            .expect("ok");
        let count_rest = ledger2
            .read_windowed_count(EntityId::new(1), "view", Window::AllTime)
            .expect("ok");
        assert_eq!(count_orig, count_rest);
        assert_eq!(count_rest, 100);
    }

    #[test]
    fn checkpoint_overwrites_previous() {
        let schema = test_schema();
        let ledger = SignalLedger::new(schema.clone(), Box::new(NoopWalWriter));
        let storage = InMemoryBackend::new();
        let ts = Timestamp::now();

        // First checkpoint: 5 entities.
        for i in 0..5u64 {
            ledger
                .record_signal("view", EntityId::new(i + 1), 1.0, ts)
                .expect("record");
        }
        ledger
            .checkpoint(
                &storage,
                CheckpointMeta {
                    checkpoint_time_ns: 1,
                    wal_sequence: 10,
                    payload_hash: [0u8; 32],
                },
            )
            .expect("checkpoint 1");

        // Add 3 more entities, then second checkpoint: 8 entities total.
        for i in 5..8u64 {
            ledger
                .record_signal("view", EntityId::new(i + 1), 1.0, ts)
                .expect("record");
        }
        ledger
            .checkpoint(
                &storage,
                CheckpointMeta {
                    checkpoint_time_ns: 2,
                    wal_sequence: 20,
                    payload_hash: [0u8; 32],
                },
            )
            .expect("checkpoint 2");

        let ledger2 = SignalLedger::new(schema, Box::new(NoopWalWriter));
        let restored_meta = ledger2
            .restore(&storage)
            .expect("restore")
            .expect("some meta");
        assert_eq!(restored_meta.wal_sequence, 20);
        assert_eq!(ledger2.entry_count(), 8);
    }

    #[test]
    fn restore_rejects_corrupted_checkpoint() {
        let schema = test_schema();
        let ledger = SignalLedger::new(schema.clone(), Box::new(NoopWalWriter));

        let now = Timestamp::now();
        for i in 0..5u64 {
            ledger
                .record_signal("view", EntityId::new(i + 1), 1.0, now)
                .expect("record");
        }

        let storage = InMemoryBackend::new();
        let meta = CheckpointMeta {
            checkpoint_time_ns: now.as_nanos(),
            wal_sequence: 100,
            payload_hash: [0u8; 32],
        };
        ledger.checkpoint(&storage, meta).expect("checkpoint");

        // Corrupt one entry value in storage.
        // Find a non-meta Sig key and overwrite its value with garbage.
        let mut corrupted = false;
        for item in storage.scan_prefix(&[]) {
            let (key, _value) = item.expect("scan");
            if let Some((entity_id, Tag::Sig, suffix)) = parse_key(&key) {
                if entity_id == EntityId::new(0) && suffix == META_SUFFIX {
                    continue;
                }
                // Overwrite with garbage bytes (same length to avoid format errors).
                let garbage = vec![0xFFu8; 983];
                let mut wb = WriteBatch::new();
                wb.put(key, garbage);
                storage.write_batch(wb).expect("write");
                corrupted = true;
                break;
            }
        }
        assert!(corrupted, "should have found an entry to corrupt");

        // Restore into a fresh ledger: should return None (integrity failure).
        let ledger2 = SignalLedger::new(schema, Box::new(NoopWalWriter));
        let result = ledger2.restore(&storage).expect("restore succeeds");
        assert!(
            result.is_none(),
            "corrupted checkpoint should return None (integrity failure)"
        );
        assert_eq!(
            ledger2.entry_count(),
            0,
            "no entries should be inserted on integrity failure"
        );
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod proptests {
    use std::time::Duration;

    use proptest::prelude::*;

    use super::*;
    use crate::schema::{DecaySpec, EntityKind, SchemaBuilder, Timestamp, Window};
    use crate::signals::ledger::NoopWalWriter;
    use crate::storage::InMemoryBackend;

    fn test_schema() -> crate::schema::Schema {
        let mut builder = SchemaBuilder::new();
        let _ = builder
            .signal(
                "view",
                EntityKind::Item,
                DecaySpec::Exponential {
                    half_life: Duration::from_secs(7 * 24 * 3600),
                },
            )
            .windows(&[Window::AllTime])
            .velocity(false)
            .add();
        builder.build().expect("valid schema")
    }

    // Full checkpoint-restore roundtrip.
    proptest! {
        #[test]
        fn checkpoint_restore_roundtrip(
            entity_count in 1usize..50,
            signals_per_entity in 1usize..20,
        ) {
            let schema = test_schema();
            let ledger = SignalLedger::new(schema.clone(), Box::new(NoopWalWriter));

            let base_ns = 1_000_000_000_000u64;
            for entity in 0..entity_count as u64 {
                for i in 0..signals_per_entity {
                    let ts = Timestamp::from_nanos(base_ns + (i as u64) * 1_000_000_000);
                    ledger
                        .record_signal("view", EntityId::new(entity + 1), 1.0, ts)
                        .unwrap();
                }
            }

            let storage = InMemoryBackend::new();
            let meta = CheckpointMeta {
                checkpoint_time_ns: base_ns,
                wal_sequence: 42,
                payload_hash: [0u8; 32], // hash computed by checkpoint()
            };
            ledger.checkpoint(&storage, meta).unwrap();

            let ledger2 = SignalLedger::new(schema, Box::new(NoopWalWriter));
            let restored_meta = ledger2.restore(&storage).unwrap();

            // The restored meta will have the computed hash, not [0u8; 32].
            // Verify timestamp and wal_sequence match.
            let restored = restored_meta.expect("should have meta");
            prop_assert_eq!(restored.checkpoint_time_ns, meta.checkpoint_time_ns);
            prop_assert_eq!(restored.wal_sequence, meta.wal_sequence);
            prop_assert_ne!(restored.payload_hash, [0u8; 32], "hash should be non-zero for non-empty checkpoints");
            prop_assert_eq!(ledger2.entry_count(), ledger.entry_count());

            for entity in 0..entity_count as u64 {
                let eid = EntityId::new(entity + 1);

                let orig_count = ledger
                    .read_windowed_count(eid, "view", Window::AllTime)
                    .unwrap();
                let rest_count = ledger2
                    .read_windowed_count(eid, "view", Window::AllTime)
                    .unwrap();
                prop_assert_eq!(orig_count, rest_count, "entity {}: all-time count mismatch", entity);
            }
        }
    }
}
