//! Checkpoint and restore for `CohortSignalLedger`.
//!
//! # Checkpoint
//!
//! `CohortSignalLedger::checkpoint()` serializes all in-memory cohort signal
//! state to the `StorageEngine` as a single atomic `WriteBatch`. Each entry
//! is serialized using the same 983-byte fixed-format record as the global
//! signal ledger (via `signals::checkpoint::format`), but keyed under
//! `Tag::Cohort` instead of `Tag::Sig`.
//!
//! # Restore
//!
//! `CohortSignalLedger::restore()` scans the storage for `Tag::Cohort` keys,
//! deserializes each entry, and populates the `DashMap`.
//!
//! # Key format
//!
//! Cohort checkpoint keys extend the standard entity-tag layout with a
//! variable-length cohort name prefix in the suffix:
//!
//! ```text
//! [entity_id: 8 BE][0x00][Tag::Cohort (0x08)][signal_type_id: 2 BE][cohort_name_len: 2 BE][cohort_name: N bytes]
//! ```
//!
//! The meta key uses `entity_id = 0` and suffix `b"meta"`.

use crate::schema::{EntityId, TidalError};
use crate::signals::SignalTypeId;
use crate::signals::checkpoint::CheckpointMeta;
use crate::signals::checkpoint::format::{deserialize_entry, serialize_entry};
use crate::signals::checkpoint::meta::{META_SUFFIX, deserialize_meta, serialize_meta};
use crate::storage::{StorageEngine, Tag, WriteBatch, encode_key, parse_key};

use super::CohortSignalLedger;

impl CohortSignalLedger {
    /// Write all in-memory cohort signal state to the storage engine atomically.
    ///
    /// Each entry is serialized as a 983-byte record (same format as the global
    /// signal ledger) and keyed under `Tag::Cohort` with a suffix encoding the
    /// signal type ID and cohort name.
    ///
    /// # Errors
    ///
    /// Returns `TidalError::Storage` if the `WriteBatch` commit or `flush` fails.
    pub fn checkpoint(
        &self,
        storage: &dyn StorageEngine,
        meta: CheckpointMeta,
    ) -> crate::Result<()> {
        let mut batch = WriteBatch::new();

        // Write checkpoint metadata at the well-known meta key.
        let meta_key = encode_key(EntityId::new(0), Tag::Cohort, META_SUFFIX);
        batch.put(meta_key, serialize_meta(&meta));

        // Write all cohort-entity-signal entries.
        for entry_ref in self.iter_entries() {
            let (cohort_name, entity_id, signal_type_id) = entry_ref.key();
            let entry = entry_ref.value();

            let suffix = cohort_entry_suffix(*signal_type_id, cohort_name);
            let key = encode_key(*entity_id, Tag::Cohort, &suffix);
            let value = serialize_entry(*entity_id, *signal_type_id, entry);
            batch.put(key, value);
        }

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

    /// Restore in-memory cohort signal state from the storage engine.
    ///
    /// Scans all keys, filters for `Tag::Cohort` entries (excluding the meta key),
    /// deserializes each entry, and inserts it into the `DashMap`.
    ///
    /// Returns `Some(CheckpointMeta)` if a checkpoint exists, or `None` on
    /// first boot (empty storage).
    ///
    /// # 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::Cohort, META_SUFFIX);
        let meta = match storage.get(&meta_key)? {
            None => None,
            Some(meta_bytes) => Some(deserialize_meta(&meta_bytes).map_err(|e| {
                TidalError::internal(
                    "cohort_checkpoint_restore",
                    format!("corrupt cohort checkpoint meta: {e}"),
                )
            })?),
        };

        // Scan all keys; keep only Tag::Cohort entry keys.
        for item in storage.scan_prefix(&[]) {
            let (key, value) = item?;
            if let Some((entity_id, Tag::Cohort, suffix)) = parse_key(&key) {
                // Skip the meta key (entity_id=0, suffix=b"meta").
                if entity_id == EntityId::new(0) && suffix == META_SUFFIX {
                    continue;
                }

                // Parse the suffix to extract signal_type_id and cohort_name.
                let Some((signal_type_id, cohort_name)) = parse_cohort_entry_suffix(suffix) else {
                    tracing::warn!(
                        entity_id = entity_id.as_u64(),
                        suffix_len = suffix.len(),
                        "skipping cohort checkpoint entry with malformed suffix"
                    );
                    continue;
                };

                let (_eid, _stid, entry) = deserialize_entry(&value).map_err(|e| {
                    TidalError::internal(
                        "cohort_checkpoint_restore",
                        format!("corrupt cohort checkpoint entry: {e}"),
                    )
                })?;

                self.insert_entry(cohort_name, entity_id, signal_type_id, entry);
            }
        }

        Ok(meta)
    }
}

// ── Suffix encoding ─────────────────────────────────────────────────────────

/// Encode a cohort entry suffix: `[signal_type_id: 2 BE][cohort_name_len: 2 BE][cohort_name: N bytes]`.
fn cohort_entry_suffix(signal_type_id: SignalTypeId, cohort_name: &str) -> Vec<u8> {
    let name_bytes = cohort_name.as_bytes();
    #[allow(clippy::cast_possible_truncation)]
    let name_len = name_bytes.len() as u16;
    let mut suffix = Vec::with_capacity(2 + 2 + name_bytes.len());
    suffix.extend_from_slice(&signal_type_id.as_u16().to_be_bytes());
    suffix.extend_from_slice(&name_len.to_be_bytes());
    suffix.extend_from_slice(name_bytes);
    suffix
}

/// Parse a cohort entry suffix back into `(SignalTypeId, cohort_name)`.
fn parse_cohort_entry_suffix(suffix: &[u8]) -> Option<(SignalTypeId, String)> {
    // Minimum: 2 (signal_type_id) + 2 (name_len) = 4 bytes
    if suffix.len() < 4 {
        return None;
    }
    let signal_type_id = SignalTypeId::new(u16::from_be_bytes([suffix[0], suffix[1]]));
    let name_len = u16::from_be_bytes([suffix[2], suffix[3]]) as usize;
    if suffix.len() < 4 + name_len {
        return None;
    }
    let name = std::str::from_utf8(&suffix[4..4 + name_len]).ok()?;
    Some((signal_type_id, name.to_owned()))
}

// ── 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::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();
        let _ = builder
            .signal(
                "like",
                EntityKind::Item,
                DecaySpec::Exponential {
                    half_life: Duration::from_secs(14 * 24 * 3600),
                },
            )
            .windows(&[Window::AllTime])
            .velocity(false)
            .add();
        builder.build().expect("valid test schema")
    }

    #[test]
    fn checkpoint_to_empty_storage() {
        let schema = test_schema();
        let ledger = CohortSignalLedger::new(&schema);

        let now = Timestamp::now();
        let type_id = ledger.resolve_signal_type("view").unwrap();
        for i in 1..=5u64 {
            ledger.record("tech", EntityId::new(i), type_id, 1.0, now.as_nanos());
        }
        ledger.record("sports", EntityId::new(1), type_id, 1.0, now.as_nanos());

        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");

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

    #[test]
    fn restore_from_empty_storage() {
        let schema = test_schema();
        let ledger = CohortSignalLedger::new(&schema);

        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 checkpoint_restore_roundtrip() {
        let schema = test_schema();
        let ledger = CohortSignalLedger::new(&schema);

        let ts_ns = Timestamp::now().as_nanos();
        let view_id = ledger.resolve_signal_type("view").unwrap();
        let like_id = ledger.resolve_signal_type("like").unwrap();

        // Record signals across cohorts.
        for i in 1..=10u64 {
            ledger.record("tech", EntityId::new(i), view_id, 1.0, ts_ns);
        }
        for i in 1..=5u64 {
            ledger.record("sports", EntityId::new(i), like_id, 2.0, ts_ns);
        }
        ledger.record("tech", EntityId::new(1), view_id, 1.0, ts_ns + 1_000_000);

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

        // Restore into a fresh ledger.
        let ledger2 = CohortSignalLedger::new(&schema);
        let restored_meta = ledger2
            .restore(&storage)
            .expect("restore")
            .expect("some meta");

        assert_eq!(restored_meta.wal_sequence, 42);
        assert_eq!(ledger2.entry_count(), ledger.entry_count());

        // Verify tech cohort windowed counts.
        let tech_count_1 = ledger2
            .read_windowed_count("tech", EntityId::new(1), "view", Window::AllTime)
            .unwrap();
        let orig_tech_count_1 = ledger
            .read_windowed_count("tech", EntityId::new(1), "view", Window::AllTime)
            .unwrap();
        assert_eq!(tech_count_1, orig_tech_count_1);
        assert_eq!(tech_count_1, 2); // Two view signals on entity 1 in tech.

        // Verify sports cohort windowed counts.
        let sports_count = ledger2
            .read_windowed_count("sports", EntityId::new(3), "like", Window::AllTime)
            .unwrap();
        assert_eq!(sports_count, 1);
    }

    #[test]
    fn checkpoint_overwrites_previous() {
        let schema = test_schema();
        let ledger = CohortSignalLedger::new(&schema);
        let storage = InMemoryBackend::new();
        let ts_ns = Timestamp::now().as_nanos();
        let view_id = ledger.resolve_signal_type("view").unwrap();

        // First checkpoint: 3 entries.
        for i in 1..=3u64 {
            ledger.record("cohort_a", EntityId::new(i), view_id, 1.0, ts_ns);
        }
        ledger
            .checkpoint(
                &storage,
                CheckpointMeta {
                    checkpoint_time_ns: 1,
                    wal_sequence: 10,
                    payload_hash: [0u8; 32],
                },
            )
            .expect("checkpoint 1");

        // Add 2 more entries, then second checkpoint: 5 entries total.
        for i in 4..=5u64 {
            ledger.record("cohort_a", EntityId::new(i), view_id, 1.0, ts_ns);
        }
        ledger
            .checkpoint(
                &storage,
                CheckpointMeta {
                    checkpoint_time_ns: 2,
                    wal_sequence: 20,
                    payload_hash: [0u8; 32],
                },
            )
            .expect("checkpoint 2");

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

    #[test]
    fn suffix_roundtrip() {
        let type_id = SignalTypeId::new(42);
        let cohort = "tech_enthusiasts_en_us";
        let suffix = cohort_entry_suffix(type_id, cohort);
        let (parsed_type_id, parsed_cohort) =
            parse_cohort_entry_suffix(&suffix).expect("valid suffix");
        assert_eq!(parsed_type_id, type_id);
        assert_eq!(parsed_cohort, cohort);
    }

    #[test]
    fn parse_malformed_suffix_returns_none() {
        // Too short.
        assert!(parse_cohort_entry_suffix(&[0x00]).is_none());
        assert!(parse_cohort_entry_suffix(&[0x00, 0x01]).is_none());
        assert!(parse_cohort_entry_suffix(&[0x00, 0x01, 0x00]).is_none());

        // Name length exceeds available bytes.
        assert!(parse_cohort_entry_suffix(&[0x00, 0x01, 0x00, 0x05, 0x61]).is_none());
    }
}
