//! 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 fixed-format record as the global signal
//! ledger (via `signals::checkpoint::format`, currently the V2 1116-byte
//! layout — `format::ENTRY_SIZE`), 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 super::CohortSignalLedger;
use crate::{
    cohort::types::{CohortName, MAX_COHORT_NAME_LEN},
    schema::{EntityId, TidalError},
    signals::{
        SignalTypeId,
        checkpoint::{
            CheckpointMeta,
            format::{deserialize_entry, serialize_entry},
            meta::{META_SUFFIX, deserialize_meta, serialize_meta},
        },
    },
    storage::{StorageEngine, Tag, WriteBatch, encode_key, parse_key},
};

impl CohortSignalLedger {
    /// Write all in-memory cohort signal state to the storage engine atomically.
    ///
    /// Each entry is serialized with the global signal ledger's fixed-format
    /// record (currently the V2 1116-byte layout, `format::ENTRY_SIZE`) and
    /// keyed under `Tag::Cohort` with a suffix encoding the signal type ID and
    /// cohort name.
    ///
    /// # Errors
    ///
    /// - `TidalError::Storage` if the `WriteBatch` commit or `flush` fails.
    /// - `TidalError::InvalidInput` if any in-memory cohort name exceeds
    ///   [`MAX_COHORT_NAME_LEN`] bytes. The suffix encodes the name length in a
    ///   `u16`, so a longer name cannot round-trip; we reject rather than
    ///   silently truncate (a truncated length corrupts the checkpoint record).
    ///   In practice `CohortRegistry::define` already rejects over-long names at
    ///   define time, so this is a defense-in-depth guard for ad-hoc/restored
    ///   cohort names that bypassed the registry.
    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.
        let mut entry_count = 0usize;
        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);
            entry_count += 1;
        }

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

        tracing::debug!(
            entry_count,
            wal_sequence = meta.wal_sequence,
            "cohort signal checkpoint written"
        );
        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).
    ///
    /// # Corruption policy
    ///
    /// The cohort aggregate is *derived* state (named user segments' decaying
    /// signal counts). It cannot be rebuilt from the WAL — this checkpoint is its
    /// only durability mechanism — so a *partial* restore would silently resurrect
    /// a wrong aggregate. We therefore treat the checkpoint as ALL-OR-NOTHING: if
    /// any structurally invalid `Tag::Cohort` record is encountered (corrupt meta,
    /// malformed key suffix, or un-deserializable value), we log a
    /// `tracing::error!` naming the corrupt record and start from an EMPTY cohort
    /// ledger (`Ok(None)`), committing none of the buffered entries.
    ///
    /// We deliberately do NOT abort the whole database open on cohort-checkpoint
    /// corruption: cohort attribution is a non-critical ranking input that
    /// re-accumulates from live signal traffic, so refusing to open the entire DB
    /// for one corrupt derived aggregate is a worse 3am outcome than serving fresh
    /// cohort state with a loud error in the log. The escalation is the `error!`
    /// record (an operator alarm), not a process abort — and the caller in
    /// `TidalDb::from_parts` no longer silently downgrades an abort to a `warn!`:
    /// code and this contract now agree. (The sibling `CommunityLedger::restore`
    /// shares the "not WAL-rebuildable" property and the same start-fresh-on-
    /// corruption posture.)
    ///
    /// # Errors
    ///
    /// - `TidalError::Storage` on a genuine I/O failure reading the checkpoint.
    ///   Checkpoint *corruption* is NOT returned as an error: it is logged at
    ///   `error!` and yields `Ok(None)` (start empty).
    pub fn restore(&self, storage: &dyn StorageEngine) -> crate::Result<Option<CheckpointMeta>> {
        // Read checkpoint metadata first. Corrupt meta -> start empty (loud).
        let meta_key = encode_key(EntityId::new(0), Tag::Cohort, META_SUFFIX);
        let meta = match storage.get(&meta_key)? {
            None => None,
            Some(meta_bytes) => match deserialize_meta(&meta_bytes) {
                Ok(m) => Some(m),
                Err(e) => {
                    tracing::error!(
                        error = %e,
                        "corrupt cohort checkpoint meta; starting from EMPTY cohort \
                         state (derived aggregate re-accumulates from live signals)"
                    );
                    return Ok(None);
                }
            },
        };

        // Buffer every entry and commit only if the ENTIRE checkpoint is intact:
        // a partial cohort aggregate is a wrong aggregate, so any corruption
        // discards everything and starts empty (see the Corruption policy above).
        //
        // TECH DEBT: scan_prefix(&[]) iterates the entire keyspace. This mirrors
        // the global signal-ledger restore scan (signals::checkpoint::restore) and
        // is correct but not selective: it touches every key, not just Tag::Cohort
        // rows. It must be replaced with a Tag::Cohort-scoped scan (e.g.
        // `scan_tag(Tag::Cohort)`) once the storage engine exposes a tag-bounded
        // iterator, so reopen cost scales with cohort state rather than total
        // keyspace size. Documented here for parity with the global-ledger note.
        let mut buffered = Vec::new();
        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;
                }

                // A malformed suffix on a Tag::Cohort key is checkpoint corruption,
                // not a foreign key. Start empty (loud) rather than commit a wrong
                // partial aggregate (see the Corruption policy section above).
                let Some((signal_type_id, cohort_name)) = parse_cohort_entry_suffix(suffix) else {
                    tracing::error!(
                        entity_id = entity_id.as_u64(),
                        suffix_len = suffix.len(),
                        "corrupt cohort checkpoint entry (malformed suffix); \
                         starting from EMPTY cohort state"
                    );
                    return Ok(None);
                };

                let entry = match deserialize_entry(&value) {
                    Ok((_eid, _stid, entry)) => entry,
                    Err(e) => {
                        tracing::error!(
                            error = %e,
                            entity_id = entity_id.as_u64(),
                            "corrupt cohort checkpoint entry (bad value); \
                             starting from EMPTY cohort state"
                        );
                        return Ok(None);
                    }
                };

                buffered.push((cohort_name, entity_id, signal_type_id, entry));
            }
        }

        // Whole checkpoint validated: commit every buffered entry.
        let restored_count = buffered.len();
        for (cohort_name, entity_id, signal_type_id, entry) in buffered {
            self.insert_entry(cohort_name, entity_id, signal_type_id, entry);
        }

        tracing::info!(
            restored_count,
            wal_sequence = meta.as_ref().map_or(0, |m| m.wal_sequence),
            "cohort signal ledger restored from checkpoint"
        );
        Ok(meta)
    }
}

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

/// Encode a cohort entry suffix: `[signal_type_id: 2 BE][cohort_name_len: 2 BE][cohort_name: N bytes]`.
///
/// # Errors
///
/// Returns `TidalError::InvalidInput` if `cohort_name` is longer than
/// [`MAX_COHORT_NAME_LEN`] bytes. The name length is stored in a `u16`, so a
/// longer name would silently truncate the length prefix and corrupt the
/// checkpoint record (and every subsequent suffix when decoded). We reject
/// instead of truncating — a checkpoint must be a faithful, round-trippable
/// image of the in-memory state.
fn cohort_entry_suffix(
    signal_type_id: SignalTypeId,
    cohort_name: &str,
) -> Result<Vec<u8>, TidalError> {
    let name_bytes = cohort_name.as_bytes();
    if name_bytes.len() > MAX_COHORT_NAME_LEN {
        return Err(TidalError::invalid_input(format!(
            "cohort name is {} bytes, exceeds maximum of {MAX_COHORT_NAME_LEN} bytes \
             (would truncate the checkpoint length prefix)",
            name_bytes.len(),
        )));
    }
    // Cast is lossless: len <= MAX_COHORT_NAME_LEN <= u16::MAX, checked above.
    #[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);
    Ok(suffix)
}

/// Parse a cohort entry suffix back into `(SignalTypeId, cohort_name)`.
///
/// The cohort name is returned as a [`CohortName`] (`Arc<str>`) so it can be
/// inserted into the ledger's `Arc<str>`-keyed map without a second copy.
fn parse_cohort_entry_suffix(suffix: &[u8]) -> Option<(SignalTypeId, CohortName)> {
    // 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;
    // Require an EXACT length match: the suffix is `[id:2][len:2][name:len]` with
    // no trailing field, so anything past `4 + name_len` is corruption, not a
    // longer name. Accepting trailing bytes would silently round-trip a record
    // whose `cohort_entry_suffix` encoder could never have produced it, masking
    // checkpoint corruption that the all-or-nothing restore policy must reject.
    if suffix.len() != 4 + name_len {
        return None;
    }
    let name = std::str::from_utf8(&suffix[4..4 + name_len]).ok()?;
    Some((signal_type_id, CohortName::from(name)))
}

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

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

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

    /// Build a `CohortName` (`Arc<str>`) for `CohortSignalLedger::record`, which
    /// keys by `&CohortName` (the production fan-out clones an `Arc` rather than
    /// allocating a `String` per matched cohort).
    fn cn(name: &str) -> CohortName {
        CohortName::from(name)
    }

    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(&cn("tech"), EntityId::new(i), type_id, 1.0, now.as_nanos());
        }
        ledger.record(
            &cn("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(&cn("tech"), EntityId::new(i), view_id, 1.0, ts_ns);
        }
        for i in 1..=5u64 {
            ledger.record(&cn("sports"), EntityId::new(i), like_id, 2.0, ts_ns);
        }
        ledger.record(
            &cn("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(&cn("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(&cn("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).expect("valid suffix");
        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 cohort_entry_suffix_rejects_over_long_name() {
        let type_id = SignalTypeId::new(1);
        // One byte over the limit must be rejected, not truncated.
        let too_long = "x".repeat(MAX_COHORT_NAME_LEN + 1);
        let err =
            cohort_entry_suffix(type_id, &too_long).expect_err("over-long name must be rejected");
        assert!(matches!(err, TidalError::InvalidInput(_)));

        // Exactly at the limit must succeed and round-trip.
        let at_limit = "y".repeat(MAX_COHORT_NAME_LEN);
        let suffix = cohort_entry_suffix(type_id, &at_limit).expect("at-limit name is valid");
        let (_id, parsed) = parse_cohort_entry_suffix(&suffix).expect("valid suffix");
        assert_eq!(&*parsed, at_limit.as_str());
    }

    #[test]
    fn checkpoint_rejects_over_long_cohort_name() {
        // A cohort name that bypassed registry validation (e.g. restored from an
        // older format) must not silently corrupt the checkpoint via truncation.
        let schema = test_schema();
        let ledger = CohortSignalLedger::new(&schema);
        let type_id = ledger.resolve_signal_type("view").unwrap();
        let now = Timestamp::now();
        let over_long = cn(&"z".repeat(MAX_COHORT_NAME_LEN + 1));
        ledger.record(&over_long, 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: 1,
            payload_hash: [0u8; 32],
        };
        let err = ledger
            .checkpoint(&storage, meta)
            .expect_err("checkpoint with over-long name must fail loudly");
        assert!(matches!(err, TidalError::InvalidInput(_)));
    }

    #[test]
    fn restore_starts_empty_on_malformed_suffix() {
        // A Tag::Cohort key with a too-short suffix is checkpoint corruption.
        // Restore must start from an EMPTY ledger (committing nothing), not commit
        // a wrong partial aggregate and not abort the whole DB open. The escalation
        // is the error! log; here we assert the all-or-nothing outcome.
        let storage = InMemoryBackend::new();

        // Write a valid meta key so restore proceeds past the meta read.
        let meta_key = encode_key(EntityId::new(0), Tag::Cohort, META_SUFFIX);
        let meta = CheckpointMeta {
            checkpoint_time_ns: 1,
            wal_sequence: 1,
            payload_hash: [0u8; 32],
        };
        storage
            .write_batch({
                let mut b = WriteBatch::new();
                b.put(meta_key, serialize_meta(&meta));
                // Malformed suffix: only 2 bytes, parse_cohort_entry_suffix needs >= 4.
                let bad_key = encode_key(EntityId::new(7), Tag::Cohort, &[0x00, 0x01]);
                b.put(bad_key, vec![0u8; 8]);
                b
            })
            .expect("write batch");

        let schema = test_schema();
        let ledger = CohortSignalLedger::new(&schema);
        let meta = ledger
            .restore(&storage)
            .expect("corruption is not an error: restore must return Ok and start empty");
        assert!(
            meta.is_none(),
            "malformed suffix must yield Ok(None) (start empty)"
        );
        assert_eq!(
            ledger.entry_count(),
            0,
            "no entries may be committed from a corrupt checkpoint"
        );
    }

    #[test]
    fn restore_starts_empty_on_corrupt_value() {
        // A Tag::Cohort key with a valid suffix but a truncated/garbage value
        // is checkpoint corruption: restore starts empty (all-or-nothing).
        let storage = InMemoryBackend::new();

        let meta_key = encode_key(EntityId::new(0), Tag::Cohort, META_SUFFIX);
        let meta = CheckpointMeta {
            checkpoint_time_ns: 1,
            wal_sequence: 1,
            payload_hash: [0u8; 32],
        };
        let good_suffix = cohort_entry_suffix(SignalTypeId::new(0), "tech").expect("valid suffix");
        storage
            .write_batch({
                let mut b = WriteBatch::new();
                b.put(meta_key, serialize_meta(&meta));
                let key = encode_key(EntityId::new(7), Tag::Cohort, &good_suffix);
                // Value far shorter than the fixed-format entry record.
                b.put(key, vec![0xAB; 3]);
                b
            })
            .expect("write batch");

        let schema = test_schema();
        let ledger = CohortSignalLedger::new(&schema);
        let meta = ledger
            .restore(&storage)
            .expect("corruption is not an error: restore must return Ok and start empty");
        assert!(
            meta.is_none(),
            "corrupt value must yield Ok(None) (start empty)"
        );
        assert_eq!(
            ledger.entry_count(),
            0,
            "no entries may be committed from a corrupt checkpoint"
        );
    }

    #[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());
    }

    #[test]
    fn parse_rejects_trailing_bytes_after_name() {
        // A well-formed suffix for signal_type_id=1, name="ab" (len=2).
        let mut suffix = vec![0x00, 0x01, 0x00, 0x02, b'a', b'b'];
        // Exact match parses.
        let (id, name) = parse_cohort_entry_suffix(&suffix).expect("exact-length suffix parses");
        assert_eq!(id, SignalTypeId::new(1));
        assert_eq!(&*name, "ab");

        // One trailing byte past the declared name length is corruption: the
        // encoder never emits trailing data, so this must be rejected rather than
        // silently round-tripped (which would mask checkpoint corruption).
        suffix.push(0xFF);
        assert!(
            parse_cohort_entry_suffix(&suffix).is_none(),
            "trailing bytes after the cohort name must be rejected"
        );
    }
}
