sol/src/brain/conversation.rs

use std::collections::{HashMap, VecDeque};

#[derive(Debug, Clone)]
pub struct ContextMessage {
    pub sender: String,
    pub content: String,
    pub timestamp: i64,
}

struct RoomContext {
    messages: VecDeque<ContextMessage>,
    max_size: usize,
}

impl RoomContext {
    fn new(max_size: usize) -> Self {
        Self {
            messages: VecDeque::with_capacity(max_size),
            max_size,
        }
    }

    fn add(&mut self, msg: ContextMessage) {
        if self.messages.len() >= self.max_size {
            self.messages.pop_front();
        }
        self.messages.push_back(msg);
    }

    fn get(&self) -> Vec<ContextMessage> {
        self.messages.iter().cloned().collect()
    }
}

pub struct ConversationManager {
    rooms: HashMap<String, RoomContext>,
    room_window: usize,
    dm_window: usize,
    max_rooms: usize,
}

impl ConversationManager {
    pub fn new(room_window: usize, dm_window: usize) -> Self {
        Self {
            rooms: HashMap::new(),
            room_window,
            dm_window,
            max_rooms: 500, // todo(sienna): make this configurable
        }
    }

    pub fn add_message(&mut self, room_id: &str, is_dm: bool, msg: ContextMessage) {
        let window = if is_dm {
            self.dm_window
        } else {
            self.room_window
        };

        // Evict oldest room if at capacity
        if !self.rooms.contains_key(room_id) && self.rooms.len() >= self.max_rooms {
            // Remove the room with the oldest latest message
            let oldest = self
                .rooms
                .iter()
                .min_by_key(|(_, ctx)| ctx.messages.back().map(|m| m.timestamp).unwrap_or(0))
                .map(|(k, _)| k.to_owned());
            if let Some(key) = oldest {
                self.rooms.remove(&key);
            }
        }

        let ctx = self
            .rooms
            .entry(room_id.to_owned())
            .or_insert_with(|| RoomContext::new(window));
        ctx.add(msg);
    }

    pub fn get_context(&self, room_id: &str) -> Vec<ContextMessage> {
        self.rooms
            .get(room_id)
            .map(|ctx| ctx.get())
            .unwrap_or_default()
    }

    pub fn room_count(&self) -> usize {
        self.rooms.len()
    }
}

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

    fn msg(sender: &str, content: &str, ts: i64) -> ContextMessage {
        ContextMessage {
            sender: sender.to_string(),
            content: content.to_string(),
            timestamp: ts,
        }
    }

    #[test]
    fn test_add_and_get_messages() {
        let mut cm = ConversationManager::new(5, 10);
        cm.add_message("!room1:x", false, msg("alice", "hello", 1));
        cm.add_message("!room1:x", false, msg("bob", "hi", 2));

        let ctx = cm.get_context("!room1:x");
        assert_eq!(ctx.len(), 2);
        assert_eq!(ctx[0].sender, "alice");
        assert_eq!(ctx[1].sender, "bob");
    }

    #[test]
    fn test_empty_room_returns_empty() {
        let cm = ConversationManager::new(5, 10);
        let ctx = cm.get_context("!nonexistent:x");
        assert!(ctx.is_empty());
    }

    #[test]
    fn test_sliding_window_group_room() {
        let mut cm = ConversationManager::new(3, 10);
        for i in 0..5 {
            cm.add_message("!room:x", false, msg("user", &format!("msg{i}"), i));
        }
        let ctx = cm.get_context("!room:x");
        assert_eq!(ctx.len(), 3);
        // Should keep the last 3
        assert_eq!(ctx[0].content, "msg2");
        assert_eq!(ctx[1].content, "msg3");
        assert_eq!(ctx[2].content, "msg4");
    }

    #[test]
    fn test_sliding_window_dm_room() {
        let mut cm = ConversationManager::new(3, 5);
        for i in 0..7 {
            cm.add_message("!dm:x", true, msg("user", &format!("dm{i}"), i));
        }
        let ctx = cm.get_context("!dm:x");
        assert_eq!(ctx.len(), 5);
        assert_eq!(ctx[0].content, "dm2");
        assert_eq!(ctx[4].content, "dm6");
    }

    #[test]
    fn test_multiple_rooms_independent() {
        let mut cm = ConversationManager::new(5, 10);
        cm.add_message("!a:x", false, msg("alice", "in room a", 1));
        cm.add_message("!b:x", false, msg("bob", "in room b", 2));

        assert_eq!(cm.get_context("!a:x").len(), 1);
        assert_eq!(cm.get_context("!b:x").len(), 1);
        assert_eq!(cm.get_context("!a:x")[0].content, "in room a");
        assert_eq!(cm.get_context("!b:x")[0].content, "in room b");
    }

    #[test]
    fn test_lru_eviction_at_max_rooms() {
        // Create a manager with max_rooms = 500 (default), but we'll use a small one
        let mut cm = ConversationManager::new(5, 10);
        cm.max_rooms = 3;

        // Add 3 rooms
        cm.add_message("!room1:x", false, msg("a", "r1", 100));
        cm.add_message("!room2:x", false, msg("b", "r2", 200));
        cm.add_message("!room3:x", false, msg("c", "r3", 300));
        assert_eq!(cm.room_count(), 3);

        // Adding a 4th room should evict the one with oldest latest message (room1, ts=100)
        cm.add_message("!room4:x", false, msg("d", "r4", 400));
        assert_eq!(cm.room_count(), 3);
        assert!(cm.get_context("!room1:x").is_empty()); // evicted
        assert_eq!(cm.get_context("!room2:x").len(), 1);
        assert_eq!(cm.get_context("!room3:x").len(), 1);
        assert_eq!(cm.get_context("!room4:x").len(), 1);
    }

    #[test]
    fn test_existing_room_not_evicted() {
        let mut cm = ConversationManager::new(5, 10);
        cm.max_rooms = 2;

        cm.add_message("!room1:x", false, msg("a", "r1", 100));
        cm.add_message("!room2:x", false, msg("b", "r2", 200));

        // Adding to existing room should NOT trigger eviction
        cm.add_message("!room1:x", false, msg("a", "r1 again", 300));
        assert_eq!(cm.room_count(), 2);
        assert_eq!(cm.get_context("!room1:x").len(), 2);
    }

    #[test]
    fn test_message_ordering_preserved() {
        let mut cm = ConversationManager::new(10, 10);
        cm.add_message("!r:x", false, msg("a", "first", 1));
        cm.add_message("!r:x", false, msg("b", "second", 2));
        cm.add_message("!r:x", false, msg("c", "third", 3));

        let ctx = cm.get_context("!r:x");
        assert_eq!(ctx[0].timestamp, 1);
        assert_eq!(ctx[1].timestamp, 2);
        assert_eq!(ctx[2].timestamp, 3);
    }
}