//! Per-user, per-creator daily notification delivery tracking.
//!
//! `NotificationTracker` records how many notifications have been delivered
//! to each `(user, creator)` pair per calendar day (days since epoch).
//! The executor consults this state when `NotificationCaps` are set on
//! a RETRIEVE query.

use std::sync::Mutex;

use dashmap::DashMap;

/// Tracks notification deliveries per `(user, creator, day)` and per `(user, day)`.
///
/// Thread-safe via `DashMap`. Each `(user_id, creator_id, day)` triple
/// maps to the number of notifications delivered on that day. A separate
/// `(user_id, day)` map tracks the cross-request total for `max_total_per_day`
/// enforcement across multiple RETRIEVE calls.
///
/// # Memory Management
///
/// Entries accumulate over time; call [`evict_old_days`](Self::evict_old_days)
/// periodically (e.g., on startup or daily) to remove stale entries. Without
/// eviction the maps grow by O(users × creators) per day of operation.
///
/// # Concurrency
///
/// Reads ([`count_today`](Self::count_today), [`count_total_today`](Self::count_total_today))
/// are lock-free via `DashMap`. The *mutating cap-enforcement* path goes
/// through [`check_and_record`](Self::check_and_record), which serializes the
/// "read the current counts, decide, then increment" sequence under
/// [`record_lock`](Self::record_lock) so two concurrent RETRIEVE responses for
/// the same user cannot both observe "under cap" and both record, overshooting
/// the per-creator or per-day cap (TOCTOU). This lock is on the query path
/// (post-diversity), not the signal-write hot path, so the contention cost is
/// acceptable.
pub struct NotificationTracker {
    /// `(user_id, creator_id, days_since_epoch)` -> count of DISTINCT items
    /// delivered today (a cap counts items notified about, deduplicated per
    /// [`counted`](Self::counted) — re-delivering the same item the same day is
    /// not a new notification).
    delivered: DashMap<(u64, u64, u32), u32>,
    /// `(user_id, days_since_epoch)` -> total distinct items delivered today.
    ///
    /// Tracks cross-request totals for `max_total_per_day` enforcement.
    user_total: DashMap<(u64, u32), u32>,
    /// `(user_id, entity_id, days_since_epoch)` -> presence marker for items
    /// already counted today. A notification cap counts *distinct items the user
    /// was notified about*, not raw delivery attempts: re-issuing the same query
    /// (or paging back to an item) must not consume the budget again. The first
    /// delivery of an item records it here and increments the counters; a repeat
    /// for the same `(user, entity, day)` is an idempotent no-op that still
    /// passes (the item is allowed to re-appear) without inflating the cap.
    counted: DashMap<(u64, u64, u32), ()>,
    /// Serializes the combined check-then-record critical section so the
    /// per-creator and per-day caps are evaluated and incremented atomically
    /// with respect to each other and to concurrent recorders.
    record_lock: Mutex<()>,
}

impl NotificationTracker {
    /// Create a new empty tracker.
    #[must_use]
    pub fn new() -> Self {
        Self {
            delivered: DashMap::new(),
            user_total: DashMap::new(),
            counted: DashMap::new(),
            record_lock: Mutex::new(()),
        }
    }

    /// Record a delivery for `(user, creator)` on `day`.
    ///
    /// Increments both the per-creator count and the per-user total.
    ///
    /// Test-only: the production cap-enforcement path goes through
    /// [`check_and_record`](Self::check_and_record), which atomically checks the
    /// caps before incrementing. This unconditional recorder exists solely to set
    /// up tracker state in tests; gating it keeps the production surface to the
    /// one path callers actually use.
    #[cfg(test)]
    pub fn record(&self, user_id: u64, creator_id: u64, day: u32) {
        *self
            .delivered
            .entry((user_id, creator_id, day))
            .or_insert(0) += 1;
        *self.user_total.entry((user_id, day)).or_insert(0) += 1;
    }

    /// Atomically enforce the daily caps and record one delivery if both pass.
    ///
    /// Under [`record_lock`](Self::record_lock) this reads the *live* current
    /// counts (including any concurrent recorder's increments), rejects if
    /// delivering one more would meet or exceed `max_total` (per-user/day) or
    /// `max_per_creator` (per-creator/day, only checked when `creator_id` is
    /// `Some`), and otherwise increments the relevant counters. Returns `true`
    /// iff the delivery was recorded.
    ///
    /// This replaces the racy "check counts, then later call `record`" pattern:
    /// two concurrent RETRIEVE responses for the same user can no longer both
    /// see headroom and both record, which previously let the response overshoot
    /// the configured caps.
    ///
    /// Idempotent per `(user_id, entity_id, day)`: an item already delivered to
    /// this user today does not consume the budget a second time. It returns
    /// `true` (the item is still allowed to re-appear, e.g. on a feed refresh or
    /// a re-issued identical query) WITHOUT re-incrementing any counter, so the
    /// daily cap counts distinct items notified about — not raw delivery
    /// attempts. The first delivery of each item enforces the caps and records.
    ///
    /// `creator_id == None` means the item has no resolvable creator: the
    /// per-creator cap is skipped (matching the executor's bucketing rule) but
    /// the per-day total cap is still enforced and incremented.
    pub fn check_and_record(
        &self,
        user_id: u64,
        creator_id: Option<u64>,
        entity_id: u64,
        day: u32,
        max_per_creator: u32,
        max_total: u32,
    ) -> bool {
        // Poisoned only if a prior holder panicked mid-update; the counters are
        // plain integers with no broken invariant, so recover the guard rather
        // than propagate a panic into the query path.
        let _guard = self
            .record_lock
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);

        // Idempotent re-delivery: an item already counted for this user today
        // is allowed through again but does not consume the budget twice. This
        // makes a re-issued identical query (or a page revisit) not double-count
        // delivered items.
        if self.counted.contains_key(&(user_id, entity_id, day)) {
            return true;
        }

        // Per-day total cap (always enforced).
        if self.count_total_today(user_id, day) >= max_total {
            return false;
        }

        // Per-creator cap (only when the item has a creator).
        if let Some(cid) = creator_id
            && self.count_today(user_id, cid, day) >= max_per_creator
        {
            return false;
        }

        // Commit: increment under the same lock so the next caller observes it,
        // and mark the item counted so a repeat is idempotent.
        if let Some(cid) = creator_id {
            *self.delivered.entry((user_id, cid, day)).or_insert(0) += 1;
        }
        *self.user_total.entry((user_id, day)).or_insert(0) += 1;
        self.counted.insert((user_id, entity_id, day), ());
        true
    }

    /// Get today's delivery count for a user-creator pair.
    #[must_use]
    pub fn count_today(&self, user_id: u64, creator_id: u64, day: u32) -> u32 {
        self.delivered
            .get(&(user_id, creator_id, day))
            .map_or(0, |c| *c)
    }

    /// Get the total number of notifications delivered to `user_id` on `day`,
    /// across all creators. Used for cross-request `max_total_per_day` enforcement.
    #[must_use]
    pub fn count_total_today(&self, user_id: u64, day: u32) -> u32 {
        self.user_total.get(&(user_id, day)).map_or(0, |c| *c)
    }

    /// Check whether delivering one more notification to `(user, creator)` on `day`
    /// would exceed the per-creator cap.
    ///
    /// Test-only: production cap checks are folded into the atomic
    /// [`check_and_record`](Self::check_and_record) under `record_lock`; a
    /// standalone read here would be a TOCTOU window. Kept for tests that assert
    /// the cap predicate in isolation.
    #[cfg(test)]
    #[must_use]
    pub fn would_exceed_per_creator(
        &self,
        user_id: u64,
        creator_id: u64,
        day: u32,
        max_per_creator: u32,
    ) -> bool {
        self.count_today(user_id, creator_id, day) >= max_per_creator
    }

    /// Check whether the user has already reached the daily total cap (cross-request).
    ///
    /// Test-only for the same reason as
    /// [`would_exceed_per_creator`](Self::would_exceed_per_creator): the live
    /// check happens atomically inside [`check_and_record`](Self::check_and_record).
    #[cfg(test)]
    #[must_use]
    pub fn would_exceed_total(&self, user_id: u64, day: u32, max_total: u32) -> bool {
        self.count_total_today(user_id, day) >= max_total
    }

    /// Evict entries for days older than `current_day - keep_days`.
    ///
    /// Call this periodically (e.g., on startup or daily) to prevent unbounded
    /// memory growth. Entries with `day < current_day - keep_days` are removed;
    /// entries for `day >= current_day - keep_days` (i.e. the last `keep_days + 1`
    /// days including today) are retained.
    ///
    /// # Example
    ///
    /// ```ignore
    /// // Keep only today and yesterday; evict everything older.
    /// tracker.evict_old_days(today, 1);
    /// ```
    pub fn evict_old_days(&self, current_day: u32, keep_days: u32) {
        let oldest_to_keep = current_day.saturating_sub(keep_days);
        self.delivered.retain(|k, _| k.2 >= oldest_to_keep);
        self.user_total.retain(|k, _| k.1 >= oldest_to_keep);
        self.counted.retain(|k, ()| k.2 >= oldest_to_keep);
    }
}

impl Default for NotificationTracker {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;

    #[test]
    fn record_and_count() {
        let tracker = NotificationTracker::new();
        assert_eq!(tracker.count_today(1, 100, 19_700), 0);
        tracker.record(1, 100, 19_700);
        assert_eq!(tracker.count_today(1, 100, 19_700), 1);
        tracker.record(1, 100, 19_700);
        assert_eq!(tracker.count_today(1, 100, 19_700), 2);
        // Different day is independent.
        assert_eq!(tracker.count_today(1, 100, 19_701), 0);
    }

    #[test]
    fn would_exceed() {
        let tracker = NotificationTracker::new();
        assert!(!tracker.would_exceed_per_creator(1, 100, 19_700, 2));
        tracker.record(1, 100, 19_700);
        assert!(!tracker.would_exceed_per_creator(1, 100, 19_700, 2));
        tracker.record(1, 100, 19_700);
        assert!(tracker.would_exceed_per_creator(1, 100, 19_700, 2));
    }

    #[test]
    fn cross_request_total_tracking() {
        let tracker = NotificationTracker::new();
        let day = 19_700u32;

        // Initial total is zero.
        assert_eq!(tracker.count_total_today(1, day), 0);
        assert!(!tracker.would_exceed_total(1, day, 3));

        // Record from two different creators.
        tracker.record(1, 100, day);
        tracker.record(1, 200, day);
        assert_eq!(tracker.count_total_today(1, day), 2);
        assert!(!tracker.would_exceed_total(1, day, 3));

        tracker.record(1, 300, day);
        assert_eq!(tracker.count_total_today(1, day), 3);
        assert!(tracker.would_exceed_total(1, day, 3));
    }

    #[test]
    fn total_is_independent_across_days() {
        let tracker = NotificationTracker::new();
        tracker.record(1, 100, 19_700);
        tracker.record(1, 200, 19_700);
        // Day 19_701 starts fresh.
        assert_eq!(tracker.count_total_today(1, 19_701), 0);
    }

    #[test]
    fn check_and_record_enforces_caps_single_threaded() {
        let tracker = NotificationTracker::new();
        let day = 19_700u32;

        // Per-creator cap = 2, total cap = 5. Distinct entity ids so each call is
        // a new item (idempotency keys on the entity).
        assert!(tracker.check_and_record(1, Some(100), 1, day, 2, 5)); // creator100: 1, total 1
        assert!(tracker.check_and_record(1, Some(100), 2, day, 2, 5)); // creator100: 2, total 2
        assert!(
            !tracker.check_and_record(1, Some(100), 3, day, 2, 5),
            "third item from creator 100 must hit the per-creator cap"
        );
        // A different creator still has headroom under the total cap.
        assert!(tracker.check_and_record(1, Some(200), 4, day, 2, 5)); // total 3
        assert_eq!(tracker.count_total_today(1, day), 3);
        assert_eq!(tracker.count_today(1, 100, day), 2);
        assert_eq!(tracker.count_today(1, 200, day), 1);
    }

    #[test]
    fn check_and_record_is_idempotent_per_entity_per_day() {
        let tracker = NotificationTracker::new();
        let day = 19_700u32;

        // First delivery of item 500 (creator 100) records once.
        assert!(tracker.check_and_record(1, Some(100), 500, day, 2, 5));
        assert_eq!(tracker.count_total_today(1, day), 1);
        assert_eq!(tracker.count_today(1, 100, day), 1);

        // Re-delivering the SAME item the same day (e.g. a re-issued identical
        // query) is allowed through but does NOT consume the budget again.
        assert!(
            tracker.check_and_record(1, Some(100), 500, day, 2, 5),
            "an already-delivered item must still be allowed to re-appear"
        );
        assert!(tracker.check_and_record(1, Some(100), 500, day, 2, 5));
        assert_eq!(
            tracker.count_total_today(1, day),
            1,
            "re-delivery must not double-count the daily total"
        );
        assert_eq!(
            tracker.count_today(1, 100, day),
            1,
            "re-delivery must not double-count the per-creator total"
        );

        // A different item still counts, and the same item on a different day is
        // a fresh notification.
        assert!(tracker.check_and_record(1, Some(100), 501, day, 2, 5));
        assert_eq!(tracker.count_today(1, 100, day), 2);
        assert!(tracker.check_and_record(1, Some(100), 500, day + 1, 2, 5));
        assert_eq!(tracker.count_today(1, 100, day + 1), 1);
    }

    #[test]
    fn check_and_record_enforces_total_cap_for_creatorless_items() {
        let tracker = NotificationTracker::new();
        let day = 1u32;
        // No creator: per-creator cap is irrelevant, total cap still applies.
        // Distinct entity ids so each is a new item.
        assert!(tracker.check_and_record(1, None, 1, day, 0, 2));
        assert!(tracker.check_and_record(1, None, 2, day, 0, 2));
        assert!(
            !tracker.check_and_record(1, None, 3, day, 0, 2),
            "creatorless deliveries must still respect the daily total cap"
        );
        assert_eq!(tracker.count_total_today(1, day), 2);
    }

    #[test]
    fn check_and_record_does_not_overshoot_under_concurrency() {
        use std::{
            sync::{
                Arc,
                atomic::{AtomicU32, Ordering as AtomicOrdering},
            },
            thread,
        };

        // Many threads race to record for the SAME (user, creator, day) with a
        // tight cap. Exactly `max_per_creator` accepts may succeed — never more.
        const THREADS: usize = 16;
        const ATTEMPTS_PER_THREAD: usize = 64;
        const MAX_PER_CREATOR: u32 = 5;
        const MAX_TOTAL: u32 = 1_000; // high enough not to be the binding cap

        let tracker = Arc::new(NotificationTracker::new());
        let accepted = Arc::new(AtomicU32::new(0));
        let day = 42u32;

        let handles: Vec<_> = (0..THREADS)
            .map(|t| {
                let tracker = Arc::clone(&tracker);
                let accepted = Arc::clone(&accepted);
                thread::spawn(move || {
                    for i in 0..ATTEMPTS_PER_THREAD {
                        // A distinct item per attempt so the per-creator cap binds
                        // on distinct items (idempotency keys on the entity).
                        let entity = (t as u64) * 1000 + i as u64;
                        if tracker.check_and_record(
                            7,
                            Some(99),
                            entity,
                            day,
                            MAX_PER_CREATOR,
                            MAX_TOTAL,
                        ) {
                            accepted.fetch_add(1, AtomicOrdering::Relaxed);
                        }
                    }
                })
            })
            .collect();
        for h in handles {
            h.join().unwrap();
        }

        // The number of accepts and the recorded count must both equal the cap
        // exactly — no overshoot from interleaved check-then-record races.
        assert_eq!(
            accepted.load(AtomicOrdering::Relaxed),
            MAX_PER_CREATOR,
            "accepts must equal the per-creator cap exactly"
        );
        assert_eq!(
            tracker.count_today(7, 99, day),
            MAX_PER_CREATOR,
            "recorded per-creator count must not overshoot the cap"
        );
        assert_eq!(
            tracker.count_total_today(7, day),
            MAX_PER_CREATOR,
            "recorded total must match accepts"
        );
    }

    #[test]
    fn check_and_record_does_not_overshoot_total_under_concurrency() {
        use std::{
            sync::{
                Arc,
                atomic::{AtomicU32, Ordering as AtomicOrdering},
            },
            thread,
        };

        // Distinct creators per attempt so the per-creator cap never binds; the
        // per-day TOTAL cap must hold under concurrency.
        const THREADS: usize = 16;
        const MAX_TOTAL: u32 = 10;

        let tracker = Arc::new(NotificationTracker::new());
        let accepted = Arc::new(AtomicU32::new(0));
        let day = 9u32;

        let handles: Vec<_> = (0..THREADS)
            .map(|t| {
                let tracker = Arc::clone(&tracker);
                let accepted = Arc::clone(&accepted);
                thread::spawn(move || {
                    for i in 0..50u64 {
                        // Unique creator AND entity id per (thread, attempt).
                        let creator = (t as u64) * 1000 + i;
                        let entity = creator;
                        if tracker.check_and_record(3, Some(creator), entity, day, 1, MAX_TOTAL) {
                            accepted.fetch_add(1, AtomicOrdering::Relaxed);
                        }
                    }
                })
            })
            .collect();
        for h in handles {
            h.join().unwrap();
        }

        assert_eq!(accepted.load(AtomicOrdering::Relaxed), MAX_TOTAL);
        assert_eq!(tracker.count_total_today(3, day), MAX_TOTAL);
    }

    #[test]
    fn evict_old_days_removes_stale_entries() {
        let tracker = NotificationTracker::new();
        let current_day = 100u32;

        // Record on old days and current day.
        tracker.record(1, 100, 97); // 3 days ago
        tracker.record(1, 100, 98); // 2 days ago
        tracker.record(1, 100, 99); // yesterday
        tracker.record(1, 100, 100); // today

        // Evict entries older than 1 day (keep today and yesterday).
        tracker.evict_old_days(current_day, 1);

        assert_eq!(
            tracker.count_today(1, 100, 97),
            0,
            "3 days ago should be evicted"
        );
        assert_eq!(
            tracker.count_today(1, 100, 98),
            0,
            "2 days ago should be evicted"
        );
        assert_eq!(
            tracker.count_today(1, 100, 99),
            1,
            "yesterday should be retained"
        );
        assert_eq!(
            tracker.count_today(1, 100, 100),
            1,
            "today should be retained"
        );

        // Total counts should also be evicted.
        assert_eq!(tracker.count_total_today(1, 97), 0);
        assert_eq!(tracker.count_total_today(1, 99), 1);
    }

    mod proptests {
        use proptest::prelude::*;

        use super::*;

        proptest! {
            /// Cross-user isolation: recording for user A does not affect user B.
            #[test]
            fn cross_user_isolation(
                user_a in 1u64..100,
                creator in 1u64..50,
                day in 1u32..100,
                recs in 1u32..10,
            ) {
                let user_b = user_a + 1000; // guaranteed different
                let tracker = NotificationTracker::new();
                for _ in 0..recs {
                    tracker.record(user_a, creator, day);
                }
                prop_assert_eq!(tracker.count_today(user_b, creator, day), 0);
                prop_assert_eq!(tracker.count_total_today(user_b, day), 0);
            }

            /// Cross-day isolation: recording on day D does not affect day D+1.
            #[test]
            fn cross_day_isolation(
                user in 1u64..100,
                creator in 1u64..50,
                day in 1u32..100,
                recs in 1u32..10,
            ) {
                let tracker = NotificationTracker::new();
                for _ in 0..recs {
                    tracker.record(user, creator, day);
                }
                prop_assert_eq!(tracker.count_today(user, creator, day + 1), 0);
                prop_assert_eq!(tracker.count_total_today(user, day + 1), 0);
            }

            /// Monotonic growth: count_today and count_total_today only increase.
            #[test]
            fn monotonic_growth(
                user in 1u64..100,
                creator in 1u64..50,
                day in 1u32..100,
                recs in 1u32..20,
            ) {
                let tracker = NotificationTracker::new();
                for i in 1..=recs {
                    tracker.record(user, creator, day);
                    prop_assert_eq!(tracker.count_today(user, creator, day), i);
                    prop_assert_eq!(tracker.count_total_today(user, day), i);
                }
            }

            /// Per-creator counts and totals stay consistent.
            #[test]
            fn per_creator_and_total_consistency(
                user in 1u64..10,
                creator_a in 1u64..50,
                creator_b in 51u64..100,
                day in 1u32..100,
                recs_a in 1u32..5,
                recs_b in 1u32..5,
            ) {
                let tracker = NotificationTracker::new();
                for _ in 0..recs_a {
                    tracker.record(user, creator_a, day);
                }
                for _ in 0..recs_b {
                    tracker.record(user, creator_b, day);
                }
                prop_assert_eq!(tracker.count_today(user, creator_a, day), recs_a);
                prop_assert_eq!(tracker.count_today(user, creator_b, day), recs_b);
                prop_assert_eq!(
                    tracker.count_total_today(user, day),
                    recs_a + recs_b,
                    "total must equal sum of per-creator counts"
                );
            }
        }
    }
}
