//! In-memory observation store backed by a `BTreeMap` keyed by content-derived
//! `ObservationId`.
//!
//! `MemStore` implements [`Store`](crate::store::Store). It is unbounded, so it
//! never returns `StoreError::Full`; allocation is therefore required and the
//! module is gated behind the `alloc` feature.

#[cfg(feature = "alloc")]
extern crate alloc;

#[cfg(feature = "alloc")]
use alloc::collections::{BTreeMap, BTreeSet};
#[cfg(feature = "alloc")]
use alloc::vec::Vec;
#[cfg(feature = "alloc")]
use crate::observation::{Observation, ObservationId, Tick};
#[cfg(feature = "alloc")]
use crate::store::{current_age, Store, StoreError, TrimPolicy};

#[cfg(feature = "alloc")]
struct Held {
    obs: Observation,
    stored_at: Tick,
}

/// An unbounded [`Store`] backed by a `BTreeMap` keyed by content-derived
/// [`ObservationId`].
///
/// Because a `BTreeMap` has no capacity limit, `MemStore` never returns
/// `StoreError::Full`. It is the only store type that requires the `alloc`
/// feature; the [`Store`] trait, [`TrimPolicy`], and [`StoreError`] are defined
/// in [`store`](crate::store) and compile without allocation.
#[cfg(feature = "alloc")]
pub struct MemStore {
    map: BTreeMap<ObservationId, Held>,
}

#[cfg(feature = "alloc")]
impl MemStore {
    /// Create an empty store.
    pub fn new() -> Self {
        MemStore {
            map: BTreeMap::new(),
        }
    }
}

#[cfg(feature = "alloc")]
impl Default for MemStore {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(feature = "alloc")]
impl Store for MemStore {
    fn insert(&mut self, obs: Observation, now: Tick) -> Result<bool, StoreError> {
        let id = obs.id();
        if self.map.contains_key(&id) {
            // Dedupe: keep the first-stored entry and its age lineage.
            return Ok(false);
        }
        self.map.insert(id, Held { obs, stored_at: now });
        Ok(true)
    }

    fn contains(&self, id: &ObservationId) -> bool {
        self.map.contains_key(id)
    }

    fn get(&self, id: &ObservationId, now: Tick) -> Option<Observation> {
        let held = self.map.get(id)?;
        let elapsed = now.saturating_sub(held.stored_at);
        Some(held.obs.aged(elapsed))
    }

    fn len(&self) -> usize {
        self.map.len()
    }

    fn for_each_id(&self, f: &mut dyn FnMut(&ObservationId)) {
        for id in self.map.keys() {
            f(id);
        }
    }

    fn for_each_topic(
        &self,
        prefix: &[u8],
        now: Tick,
        f: &mut dyn FnMut(&Observation),
    ) {
        for held in self.map.values() {
            if held.obs.topic.starts_with(prefix) {
                let elapsed = now.saturating_sub(held.stored_at);
                f(&held.obs.aged(elapsed));
            }
        }
    }

    fn trim(&mut self, policy: &TrimPolicy, now: Tick) -> usize {
        let mut removed: BTreeSet<ObservationId> = BTreeSet::new();

        // (a) keep_per_topic: per exact topic, drop everything but the `N`
        // youngest (lowest current age).
        if let Some(keep) = policy.keep_per_topic {
            let mut groups: BTreeMap<&[u8], Vec<ObservationId>> = BTreeMap::new();
            for (id, held) in self.map.iter() {
                if !removed.contains(id) {
                    groups.entry(held.obs.topic.as_slice()).or_default().push(*id);
                }
            }
            for member_ids in groups.values() {
                if member_ids.len() <= keep {
                    continue;
                }
                let mut ids = member_ids.clone();
                ids.sort_by(|&a, &b| {
                    current_age(&self.map[&a].obs, self.map[&a].stored_at, now)
                        .cmp(&current_age(&self.map[&b].obs, self.map[&b].stored_at, now))
                });
                for id in ids.into_iter().skip(keep) {
                    removed.insert(id);
                }
            }
        }

        // (b) max_age_ms: drop anything whose current age exceeds the cutoff.
        if let Some(max_age) = policy.max_age_ms {
            for (id, held) in self.map.iter() {
                if !removed.contains(id)
                    && current_age(&held.obs, held.stored_at, now) > max_age
                {
                    removed.insert(*id);
                }
            }
        }

        // (c) max_count: drop the oldest (highest current age) until the number
        // of remaining (untrimmed) observations is at most `max_count`.
        if let Some(max_count) = policy.max_count {
            let mut remaining: Vec<ObservationId> = self
                .map
                .keys()
                .copied()
                .filter(|id| !removed.contains(id))
                .collect();
            if remaining.len() > max_count {
                remaining.sort_by(|&a, &b| {
                    current_age(&self.map[&a].obs, self.map[&a].stored_at, now)
                        .cmp(&current_age(&self.map[&b].obs, self.map[&b].stored_at, now))
                });
                let over = remaining.len() - max_count;
                for id in remaining.into_iter().rev().take(over) {
                    removed.insert(id);
                }
            }
        }

        self.map.retain(|id, _| !removed.contains(id));
        removed.len()
    }
}

#[cfg(all(test, feature = "alloc"))]
mod tests {
    use super::MemStore;
    use crate::observation::{Observation, Payload, Topic};
    use crate::store::{Store, TrimPolicy};

    fn obs(observer: [u8; 16], seq: u64, age: u64, at: Option<u64>, topic: &[u8]) -> Observation {
        Observation {
            observer,
            seq,
            age_ms: age,
            observed_at: at,
            topic: Topic::from_slice(topic).unwrap(),
            payload: Payload::from_slice(b"").unwrap(),
        }
    }

    #[test]
    fn insert_then_contains_and_len() {
        let mut s = MemStore::new();
        let o = obs([1; 16], 1, 0, None, b"door/1");
        let id = o.id();
        assert!(s.insert(o, 1000).unwrap());
        assert!(s.contains(&id));
        assert_eq!(s.len(), 1);
    }

    #[test]
    fn duplicate_id_is_not_reinserted() {
        let mut s = MemStore::new();
        let o1 = obs([1; 16], 1, 100, None, b"door/1");
        let id = o1.id();
        assert!(s.insert(o1, 1000).unwrap());
        // Same (observer, seq, topic, payload), different age -> deduped.
        let o2 = obs([1; 16], 1, 999, None, b"door/1");
        assert_eq!(id, o2.id());
        assert!(!s.insert(o2, 2000).unwrap());
        assert_eq!(s.len(), 1);
        // get() still returns the first-stored age lineage.
        let g = s.get(&id, 1500).unwrap();
        assert_eq!(g.age_ms, 600);
    }

    #[test]
    fn get_adjusts_age_by_elapsed() {
        let mut s = MemStore::new();
        let o = obs([1; 16], 1, 100, None, b"door/1");
        let id = o.id();
        s.insert(o, 1000).unwrap();
        let g = s.get(&id, 1500).unwrap();
        assert_eq!(g.age_ms, 600);
    }

    #[test]
    fn get_with_backwards_tick_saturates() {
        let mut s = MemStore::new();
        let o = obs([1; 16], 1, 100, None, b"door/1");
        let id = o.id();
        s.insert(o, 1000).unwrap();
        let g = s.get(&id, 500).unwrap();
        assert_eq!(g.age_ms, 100);
    }

    #[test]
    fn for_each_topic_prefix_matches() {
        let mut s = MemStore::new();
        s.insert(obs([1; 16], 1, 0, None, b"door/1"), 0).unwrap();
        s.insert(obs([1; 16], 2, 0, None, b"door/2"), 0).unwrap();
        s.insert(obs([1; 16], 3, 0, None, b"wall/1"), 0).unwrap();
        let mut count = 0;
        s.for_each_topic(b"door/", 0, &mut |_o| {
            count += 1;
        });
        assert_eq!(count, 2);
    }

    #[test]
    fn for_each_topic_exact_empty_prefix_is_all() {
        let mut s = MemStore::new();
        s.insert(obs([1; 16], 1, 0, None, b"door/1"), 0).unwrap();
        s.insert(obs([1; 16], 2, 0, None, b"door/2"), 0).unwrap();
        s.insert(obs([1; 16], 3, 0, None, b"wall/1"), 0).unwrap();
        let mut count = 0;
        s.for_each_topic(b"", 0, &mut |_o| {
            count += 1;
        });
        assert_eq!(count, s.len());
    }

    #[test]
    fn trim_keep_per_topic_keeps_youngest() {
        let mut s = MemStore::new();
        for (i, age) in [10u64, 20, 30, 40, 50].iter().enumerate() {
            s.insert(obs([1; 16], i as u64, *age, None, b"t"), 0).unwrap();
        }
        let policy = TrimPolicy {
            keep_per_topic: Some(2),
            max_age_ms: None,
            max_count: None,
        };
        assert_eq!(s.trim(&policy, 0), 3);
        let mut ages = std::collections::BTreeSet::new();
        s.for_each_topic(b"t", 0, &mut |o| {
            ages.insert(o.age_ms);
        });
        assert_eq!(ages, std::collections::BTreeSet::from_iter([10u64, 20]));
    }

    #[test]
    fn trim_max_age_drops_old() {
        let mut s = MemStore::new();
        s.insert(obs([1; 16], 1, 10, None, b"t"), 0).unwrap();
        s.insert(obs([1; 16], 2, 10_000, None, b"t"), 0).unwrap();
        let policy = TrimPolicy {
            keep_per_topic: None,
            max_age_ms: Some(1000),
            max_count: None,
        };
        assert_eq!(s.trim(&policy, 0), 1);
        let mut found = false;
        s.for_each_topic(b"t", 0, &mut |o| {
            assert_eq!(o.age_ms, 10);
            found = true;
        });
        assert!(found);
    }

    #[test]
    fn trim_max_count_drops_oldest() {
        let mut s = MemStore::new();
        for i in 0..10u64 {
            s.insert(obs([1; 16], i, (i + 1) * 10, None, b"t"), 0).unwrap();
        }
        let policy = TrimPolicy {
            keep_per_topic: None,
            max_age_ms: None,
            max_count: Some(4),
        };
        assert_eq!(s.trim(&policy, 0), 6);
        assert_eq!(s.len(), 4);
        let mut ages = std::collections::BTreeSet::new();
        s.for_each_topic(b"t", 0, &mut |o| {
            ages.insert(o.age_ms);
        });
        assert_eq!(ages, std::collections::BTreeSet::from_iter([10u64, 20, 30, 40]));
    }

    #[test]
    fn trimmed_obs_can_be_reinserted() {
        let mut s = MemStore::new();
        let o = obs([1; 16], 1, 10, None, b"t");
        let id = o.id();
        assert!(s.insert(o, 0).unwrap());
        assert_eq!(s.len(), 1);
        let policy = TrimPolicy {
            keep_per_topic: None,
            max_age_ms: Some(1),
            max_count: None,
        };
        assert_eq!(s.trim(&policy, 0), 1);
        assert_eq!(s.len(), 0);
        assert!(!s.contains(&id));
        // Trimmed then re-received is accepted.
        let o2 = obs([1; 16], 1, 20, None, b"t");
        assert_eq!(id, o2.id());
        assert!(s.insert(o2, 0).unwrap());
        assert!(s.contains(&id));
        assert_eq!(s.len(), 1);
    }
}