pub mod base_agent;
pub mod goal_agent;

use crate::{
    llms::{
        anthropic::AnthropicClient,
        messages::MessageType,
        open_ai::OpenAIClient,
        tools::{ToolDefinition, ToolExecution},
    },
    prompts::BasePrompt,
};
use anyhow::Result;

/**
 * A module that handles agent behavior
 *
 * What is an agent?
 * - A collection of messages that are used with subsequent calls
 * - A function to call the agent
 */

pub enum ProviderClient {
    Anthropic(AnthropicClient),
    OpenAI(OpenAIClient),
}

pub trait AgentBehavior {
    fn call(&mut self, prompt: impl BasePrompt) -> Result<Vec<MessageType>>;
}

pub trait GoalSeekingBehavior {
    /// Chase a goal for a given number of iterations
    ///
    /// Iterations should be used for LLM based agents. Any other limits should be handled
    /// within the agent itself.
    fn chase(&mut self, goal: impl BasePrompt, iterations: u32) -> Result<Vec<MessageType>>;
}

pub trait ToolUsageBehavior {
    fn use_tool(&mut self, tool: impl ToolExecution) -> Result<Vec<MessageType>>;
    fn add_tool(&mut self, tool: impl ToolDefinition);
    fn remove_tool(&mut self, tool: impl ToolDefinition);
    fn get_tools(&self) -> Vec<impl ToolDefinition>;
    fn set_tools(&mut self, tools: Vec<impl ToolDefinition>);
    fn clear_tools(&mut self);
    fn get_tool_by_name(&self, name: &str) -> Option<impl ToolDefinition>;
}

#[cfg(test)]
mod tests {
    use crate::{context::ContextSession, prompts::templates};

    use super::*;

    #[test]
    fn test_get_prompt() -> Result<()> {
        dotenvy::dotenv().ok();
        let mut context = ContextSession::new()?;
        context.add_text_to_context("test", "Hello, world!");
        let prompt = templates::get_context_prompt("Hello, world!".to_string(), &context);

        assert!(!prompt.is_empty());
        Ok(())
    }
}
