rsclaw 0.0.1-alpha.1

use crate::provider::{ChatMessage, MessageRole};
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::Arc;

/// Agent conversation context.
#[derive(Debug, Clone)]
pub struct AgentContext {
    pub system_prompt: Arc<str>,
    pub messages: VecDeque<ChatMessage>,
    pub max_messages: usize,
    pub token_count: usize,
    pub max_tokens: usize,
}

impl AgentContext {
    /// Create a new agent context.
    pub fn new(system_prompt: Arc<str>, max_messages: usize, max_tokens: usize) -> Self {
        Self {
            system_prompt,
            messages: VecDeque::new(),
            max_messages,
            token_count: 0,
            max_tokens,
        }
    }

    /// Add a user message.
    pub fn add_user_message(&mut self, content: impl Into<String>) {
        let msg = ChatMessage::user(content);
        self.token_count += estimate_tokens(&msg.content);
        self.messages.push_back(msg);
        self.trim_if_needed();
    }

    /// Add an assistant message.
    pub fn add_assistant_message(&mut self, content: impl Into<String>) {
        let msg = ChatMessage::assistant(content);
        self.token_count += estimate_tokens(&msg.content);
        self.messages.push_back(msg);
        self.trim_if_needed();
    }

    /// Get all messages for LLM request.
    pub fn get_messages(&self) -> Vec<ChatMessage> {
        let mut msgs = vec![ChatMessage::system(self.system_prompt.as_ref())];
        msgs.extend(self.messages.iter().cloned());
        msgs
    }

    /// Trim messages if exceeding limits.
    fn trim_if_needed(&mut self) {
        while self.messages.len() > self.max_messages || self.token_count > self.max_tokens {
            if let Some(msg) = self.messages.pop_front() {
                self.token_count = self
                    .token_count
                    .saturating_sub(estimate_tokens(&msg.content));
            } else {
                break;
            }
        }
    }

    /// Clear all messages.
    pub fn clear(&mut self) {
        self.messages.clear();
        self.token_count = 0;
    }

    /// Get message count.
    pub fn message_count(&self) -> usize {
        self.messages.len()
    }
}

/// Estimate token count (simple approximation).
fn estimate_tokens(text: &str) -> usize {
    text.len() / 4
}

impl Default for AgentContext {
    fn default() -> Self {
        Self::new(Arc::from("You are a helpful assistant."), 100, 4000)
    }
}