use std::fmt;

use serde::{Deserialize, Serialize};

use crate::{function_specification::FunctionSpecification, message::Message};

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ChatContext {
    pub model: String,
    pub messages: Vec<Message>,
    pub functions: Vec<FunctionSpecification>,
    pub function_call: Option<String>,
}

impl ChatContext {
    /// Creates a new ChatContext with a model name
    /// as a string. This is an internal function used by other functions.
    pub fn new(model: String) -> ChatContext {
        ChatContext {
            model,
            messages: Vec::new(),
            functions: Vec::new(),
            function_call: None,
        }
    }

    /// Pushes a message in the chat context
    /// as a Message. This is an internal function used by other functions.
    /// It is recommended to use ChatGPT.push_message()
    pub fn push_message(&mut self, message: Message) {
        self.messages.push(message);
    }

    /// Sets the messages in the chat context
    /// as a vector of Message.
    /// This is an internal function used by other functions.
    pub fn set_messages(&mut self, messages: Vec<Message>) {
        self.messages = messages;
    }

    /// Pushes a function in the chat context
    /// as a FunctionSpecification.
    /// This is an internal function used by other functions.
    /// It is recommended to use ChatGPT.push_function()
    pub fn push_function(&mut self, functions: FunctionSpecification) {
        self.functions.push(functions);
    }

    /// Sets the functions in the chat context
    /// as a vector of FunctionSpecification.
    /// This is an internal function used by other functions.
    pub fn set_functions(&mut self, functions: Vec<FunctionSpecification>) {
        self.functions = functions;
    }

    /// Sets the last message sent by the user or the bot
    /// as a string. This is an internal function used by other functions.
    pub fn set_function_call(&mut self, function_call: String) {
        self.function_call = Some(function_call);
    }

    /// Returns the last message sent by the user or the bot
    /// as a string. This is an internal function used by other functions.
    /// It is recommended to use ChatGPT.last_content()
    pub fn last_content(&self) -> Option<String> {
        match self.messages.last() {
            Some(message) => {
                if let Some(c) = message.content.clone() {
                    Some(c)
                } else {
                    None
                }
            }
            None => None,
        }
    }

    /// Returns the last function call in the chat context
    /// as a tuple of the function name and the arguments.
    /// This is an internal function used by other functions.
    /// It is recommended to use ChatGPT.last_function_call()
    pub fn last_function_call(&self) -> Option<(String, String)> {
        match self.messages.last() {
            Some(message) => {
                if let Some(f) = message.function_call.clone() {
                    Some((f.name, f.arguments))
                } else {
                    None
                }
            }
            None => None,
        }
    }
}

// Print valid JSON for ChatContext, no commas if last field
impl fmt::Display for ChatContext {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{{\"model\":\"{}\"", self.model)?;
        if !self.messages.is_empty() {
            write!(f, ",\"messages\":[")?;
            for (i, message) in self.messages.iter().enumerate() {
                write!(f, "{}", message)?;
                if i < self.messages.len() - 1 {
                    write!(f, ",")?;
                }
            }
            write!(f, "]")?;
        }
        if self.functions.len() > 0 {
            write!(f, ",\"functions\":[")?;
            for (i, function) in self.functions.iter().enumerate() {
                write!(f, "{}", function)?;
                if i < self.functions.len() - 1 {
                    write!(f, ",")?;
                }
            }
            write!(f, "]")?;
        }
        if let Some(function_call) = &self.function_call {
            write!(f, ",\"function_call\":\"{}\"", function_call)?;
        }
        write!(f, "}}")
    }
}
#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::*;
    use crate::{
        function_specification::{Parameters, Property},
        message::Message,
    };

    #[test]
    fn test_display_for_chat_context() {
        let mut chat_context = ChatContext::new("test_model".to_string());
        chat_context.push_message(Message {
            role: "role".to_string(),
            content: Some("Hello".to_string()),
            name: None,
            function_call: None,
        });
        chat_context.push_message(Message {
            role: "bot".to_string(),
            content: Some("Hi".to_string()),
            name: None,
            function_call: None,
        });
        assert_eq!(
            chat_context.to_string(),
            "{\"model\":\"test_model\",\"messages\":[{\"role\":\"role\",\"content\":\"Hello\"},{\"role\":\"bot\",\"content\":\"Hi\"}]}"
        );
    }

    #[test]
    fn test_display_chat_context_with_functions() {
        let mut chat_context = ChatContext::new("test_model".to_string());

        // Add a function to the chat context
        let mut properties = HashMap::new();
        properties.insert(
            "location".to_string(),
            Property {
                type_: "string".to_string(),
                description: Some("a dummy string".to_string()),
                enum_: None,
            },
        );
        let function = FunctionSpecification {
            name: "test_function".to_string(),
            description: Some("a dummy function to test the chat context".to_string()),
            parameters: Some(Parameters {
                type_: "object".to_string(),
                properties,
                required: vec!["location".to_string()],
            }),
        };
        chat_context.push_function(function);

        // Add a message to the chat context
        chat_context.push_message(Message {
            role: "test".to_string(),
            content: Some("hi".to_string()),
            name: Some("test_function".to_string()),
            function_call: None, // Lets assume a function has not been called yet
        });

        // Print the chat context, with the model, the messages, the functions, and the function_call
        assert_eq!(
            chat_context.to_string(),
            "{\"model\":\"test_model\",\"messages\":[{\"role\":\"test\",\"content\":\"hi\",\"name\":\"test_function\"}],\"functions\":[{\"name\":\"test_function\",\"description\":\"a dummy function to test the chat context\",\"parameters\":{\"type\":\"object\",\"properties\":{\"location\":{\"type\":\"string\",\"description\":\"a dummy string\"}},\"required\":[\"location\"]}}]}"
        );
    }

    #[test]
    fn test_last_content() {
        let mut chat_context = ChatContext::new("model".to_string());

        // Test with no messages
        assert_eq!(chat_context.last_content(), None);

        // Test with a message with no content
        chat_context.push_message(Message {
            role: "role".to_string(),
            content: None,
            name: None,
            function_call: None,
        });
        assert_eq!(chat_context.last_content(), None);

        // Test with a message with content
        chat_context.push_message(Message {
            role: "role".to_string(),
            content: Some("content".to_string()),
            name: None,
            function_call: None,
        });
        assert_eq!(chat_context.last_content(), Some("content".to_string()));
    }

    #[test]
    fn test_last_function_call() {
        let mut chat_context = ChatContext::new("model".to_string());

        // Test with no messages
        assert_eq!(chat_context.last_content(), None);

        // Test with a message with no function call
        chat_context.push_message(Message {
            role: "role".to_string(),
            content: None,
            name: None,
            function_call: None,
        });
        assert_eq!(chat_context.last_content(), None);

        // Test with a message with function call
        use crate::message::FunctionCall;
        chat_context.push_message(Message {
            role: "role".to_string(),
            content: None,
            name: None,
            function_call: Some(FunctionCall {
                name: "function".to_string(),
                arguments: "arguments".to_string(),
            }),
        });
        assert_eq!(
            chat_context.last_function_call(),
            Some(("function".to_string(), "arguments".to_string()))
        );
    }
}