use crate::{
    llms::{
        anthropic::{
            messages::{AnthropicMessage, AnthropicMessageContent},
            AnthropicClient,
        },
        messages::{AbstractMessage, Message, MessageType},
        open_ai::{messages::OpenAIMessage, OpenAIClient},
        LLMProvider,
    },
    prompts::BasePrompt,
};

use super::{AgentBehavior, ProviderClient};

use anyhow::{Ok, Result};

pub struct BaseAgent {
    client: ProviderClient,
    messages: Vec<MessageType>,
}

impl AgentBehavior for BaseAgent {
    fn call(&mut self, prompt: impl BasePrompt) -> Result<Vec<MessageType>> {
        // Add the prompt as a message. Split this out since borrowing mut client
        match &mut self.client {
            ProviderClient::Anthropic(_client) => {
                let message = prompt.to_anthropic()?;
                self.add_message(MessageType::Anthropic(message.clone()));
            }
            ProviderClient::OpenAI(_client) => {
                let message = prompt.to_open_ai()?;
                self.add_message(MessageType::OpenAI(message.clone()));
            }
        }

        // Get the response based on the provider client
        let response: Vec<MessageType> = match &mut self.client {
            ProviderClient::Anthropic(client) => {
                let messages: Vec<AnthropicMessage> =
                    Self::convert_to_anthropic_messages(&self.messages)?;

                client
                    .get_completion(messages)?
                    .into_iter()
                    .map(MessageType::Anthropic)
                    .collect()
            }
            ProviderClient::OpenAI(client) => {
                let messages: Result<Vec<OpenAIMessage>> =
                    Self::convert_to_open_ai_messages(&self.messages);
                client
                    .get_completion(messages?)?
                    .into_iter()
                    .map(MessageType::OpenAI)
                    .collect()
            }
        };

        // Add the response messages and return all messages
        self.messages.extend(response);
        Ok(self.messages.clone())
    }
}

impl BaseAgent {
    pub fn new(client: ProviderClient) -> Self {
        Self {
            client,
            messages: vec![],
        }
    }

    pub fn open_ai() -> Self {
        Self::new(ProviderClient::OpenAI(OpenAIClient::new()))
    }

    pub fn anthropic() -> Self {
        Self::new(ProviderClient::Anthropic(AnthropicClient::new()))
    }

    fn convert_to_anthropic_messages(messages: &[MessageType]) -> Result<Vec<AnthropicMessage>> {
        Ok(messages
            .iter()
            .map(|m| {
                Ok(AnthropicMessage::new(
                    vec![AnthropicMessageContent::Text {
                        text: m.get_content().map_err(|e| anyhow::anyhow!(e))?,
                    }],
                    m.get_role(),
                ))
            })
            .collect::<Result<Vec<_>, _>>()?)
    }

    fn convert_to_open_ai_messages(messages: &[MessageType]) -> Result<Vec<OpenAIMessage>> {
        messages
            .iter()
            .map(|m| -> Result<OpenAIMessage> {
                Ok(OpenAIMessage::new(
                    &m.get_content().map_err(|e| anyhow::anyhow!(e))?,
                    m.get_role(),
                ))
            })
            .collect::<Result<Vec<_>, _>>()
    }

    pub fn get_client(&self) -> &ProviderClient {
        &self.client
    }

    pub fn add_message(&mut self, message: MessageType) {
        self.messages.push(message);
    }

    pub fn get_messages(&self) -> &Vec<MessageType> {
        &self.messages
    }

    pub fn clear_messages(&mut self) {
        self.messages.clear();
    }

    pub fn get_last_message(&self) -> Option<&MessageType> {
        self.messages.last()
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use crate::context::ContextSession;
    use crate::llms::messages::Role;
    use crate::prompts::templates::get_context_prompt;
    use crate::prompts::{ProviderPrompt, UserPrompt};

    #[test]
    fn test_add_message() -> Result<()> {
        dotenvy::dotenv().ok();

        let mut context = ContextSession::new()?;
        context.add_text_to_context("test", "Hello, world!");

        let mut agent = BaseAgent::open_ai();
        agent.add_message(MessageType::OpenAI(OpenAIMessage::new(
            "Hello, world!",
            Role::User,
        )));
        Ok(())
    }

    #[test]
    fn test_add_prompt_as_message() -> Result<()> {
        dotenvy::dotenv().ok();

        let mut context = ContextSession::new()?;
        context.add_text_to_context("test", "Hello, world!");
        let mut agent = BaseAgent::open_ai();
        let prompt = UserPrompt(get_context_prompt("Hello, world!".to_string(), &context));

        agent.add_message(MessageType::OpenAI(prompt.to_open_ai()?));

        assert!(!agent.messages.last().unwrap().get_content()?.is_empty());
        Ok(())
    }
}
