use serde::{Deserialize, Serialize};

use crate::llms::messages::{AbstractMessage, Message, Role};

use anyhow::Result;

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct OpenAIFunction {
    pub name: String,
    pub arguments: String, // JSON string
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct OpenAIToolCall {
    pub id: String,
    pub r#type: String,
    pub function: OpenAIFunction,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(tag = "role", rename_all = "lowercase")]
pub enum OpenAIMessage {
    System {
        content: String,
    },
    User {
        content: String,
    },
    Assistant {
        content: Option<String>,
        tool_calls: Option<Vec<OpenAIToolCall>>,
    },
    Tool {
        content: String,
        tool_call_id: String,
    },
}

impl AbstractMessage for OpenAIMessage {
    fn get_role(&self) -> Role {
        match self {
            OpenAIMessage::System { .. } => Role::System,
            OpenAIMessage::User { .. } => Role::User,
            OpenAIMessage::Assistant { .. } => Role::Assistant,
            OpenAIMessage::Tool { .. } => Role::Tool,
        }
    }

    fn get_content(&self) -> Result<String> {
        match self {
            OpenAIMessage::System { content, .. } => Ok(content.clone()),
            OpenAIMessage::User { content, .. } => Ok(content.clone()),
            OpenAIMessage::Assistant {
                content,
                tool_calls,
                ..
            } => match content {
                Some(content) => Ok(content.clone()),
                None => match tool_calls {
                    Some(tool_calls) => Err(anyhow::anyhow!(
                        "No content, this is a tool call: {:?}",
                        tool_calls
                    )),
                    None => Err(anyhow::anyhow!("No content")),
                },
            },
            OpenAIMessage::Tool { content, .. } => Ok(content.clone()),
        }
    }
}

impl OpenAIMessage {
    pub fn user(content: &str) -> Self {
        Self::User {
            content: content.to_string(),
        }
    }

    pub fn assistant(content: &str) -> Self {
        Self::Assistant {
            content: Some(content.to_string()),
            tool_calls: None,
        }
    }

    pub fn tool(content: &str, tool_call_id: &str) -> Self {
        Self::Tool {
            content: content.to_string(),
            tool_call_id: tool_call_id.to_string(),
        }
    }
}

impl Message for OpenAIMessage {
    fn new(content: &str, role: Role) -> Self {
        match role {
            Role::System => Self::System {
                content: content.to_string(),
            },
            Role::User => Self::User {
                content: content.to_string(),
            },
            Role::Assistant => Self::Assistant {
                content: Some(content.to_string()),
                tool_calls: None,
            },
            _ => panic!("Can't create a tool message. Must be called from a tool call"),
        }
    }
}
