agentlib-memory 0.1.0

use agentlib_core::{
    MemoryProvider, MemoryReadOptions, MemoryWriteOptions, ModelMessage, Role, async_trait,
    trim_to_token_budget,
};
use anyhow::Result;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;

pub struct SlidingWindowMemory {
    max_tokens: usize,
    max_turns: usize,
    store: Arc<Mutex<HashMap<String, Vec<ConversationTurn>>>>,
}

#[derive(Debug, Clone)]
struct ConversationTurn {
    messages: Vec<ModelMessage>,
}

impl SlidingWindowMemory {
    pub fn new(max_tokens: usize, max_turns: usize) -> Self {
        Self {
            max_tokens,
            max_turns,
            store: Arc::new(Mutex::new(HashMap::new())),
        }
    }
}

#[async_trait]
impl MemoryProvider for SlidingWindowMemory {
    async fn read(&self, options: MemoryReadOptions) -> Result<Vec<ModelMessage>> {
        let session_id = options.session_id.unwrap_or_else(|| "default".to_string());
        let store = self.store.lock().await;
        let turns = store.get(&session_id).cloned().unwrap_or_default();

        // Flatten turns to messages
        let messages: Vec<ModelMessage> = turns.into_iter().flat_map(|t| t.messages).collect();

        // Apply token budget
        let trimmed = trim_to_token_budget(messages, self.max_tokens);
        Ok(trimmed)
    }

    async fn write(&self, messages: Vec<ModelMessage>, options: MemoryWriteOptions) -> Result<()> {
        let session_id = options.session_id.unwrap_or_else(|| "default".to_string());

        let non_system: Vec<ModelMessage> = messages
            .into_iter()
            .filter(|m| m.role != Role::System)
            .collect();

        let new_turns = group_into_turns(non_system);

        let mut store = self.store.lock().await;
        let existing = store.entry(session_id).or_default();

        existing.extend(new_turns);

        // Keep last max_turns
        if existing.len() > self.max_turns {
            *existing = existing.split_off(existing.len() - self.max_turns);
        }

        Ok(())
    }

    async fn clear(&self, session_id: Option<&str>) -> Result<()> {
        let mut store = self.store.lock().await;
        if let Some(sid) = session_id {
            store.remove(sid);
        } else {
            store.clear();
        }
        Ok(())
    }
}

fn group_into_turns(messages: Vec<ModelMessage>) -> Vec<ConversationTurn> {
    let mut turns = Vec::new();
    let mut current = Vec::new();

    for msg in messages {
        let role = msg.role;
        let has_tool_calls = msg
            .tool_calls
            .as_ref()
            .map(|tc| !tc.is_empty())
            .unwrap_or(false);
        current.push(msg);

        // A turn ends when the assistant finishes (no pending tool calls)
        if role == Role::Assistant && !has_tool_calls {
            turns.push(ConversationTurn { messages: current });
            current = Vec::new();
        }
    }

    if !current.is_empty() {
        turns.push(ConversationTurn { messages: current });
    }

    turns
}