use crate::llms::{
    anthropic::messages::{AnthropicMessage, AnthropicMessageContent},
    messages::MessageType,
    open_ai::messages::OpenAIMessage,
};
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ChatMessageSource {
    User,
    Assistant,
    System,
    Tool,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ChatMessageModelProvider {
    OpenAI,
    Anthropic,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ChatMessageModel {
    #[serde(rename = "gpt-4o")]
    GPT4o,
    #[serde(rename = "claude3-5-sonnet")]
    ClaudeSonnet3,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub struct Tool {
    id: String,
    name: Option<String>,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
enum ContentTypes {
    List(Vec<AnthropicMessageContent>),
    Text(String),
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub struct ChatMessage {
    role: String,
    content: Option<ContentTypes>,
    tool: Option<Tool>,
    source: ChatMessageSource,
    model: ChatMessageModel,
    model_provider: ChatMessageModelProvider,
}

impl ChatMessage {
    pub fn from_message_type(message: MessageType) -> Self {
        match message {
            MessageType::Anthropic(m) => {
                Self::from_anthropic_message(m, ChatMessageModel::ClaudeSonnet3)
            }
            MessageType::OpenAI(m) => Self::from_open_ai_message(m, ChatMessageModel::GPT4o),
        }
    }

    pub fn from_open_ai_message(message: OpenAIMessage, model: ChatMessageModel) -> Self {
        match message {
            OpenAIMessage::System { content, .. } => Self {
                role: "system".to_string(),
                content: Some(ContentTypes::Text(content)),
                tool: None,
                source: ChatMessageSource::System,
                model,
                model_provider: ChatMessageModelProvider::OpenAI,
            },
            OpenAIMessage::User { content, .. } => Self {
                role: "user".to_string(),
                content: Some(ContentTypes::Text(content)),
                tool: None,
                source: ChatMessageSource::User,
                model,
                model_provider: ChatMessageModelProvider::OpenAI,
            },
            OpenAIMessage::Assistant { content, .. } => Self {
                role: "assistant".to_string(),
                content: content.map(ContentTypes::Text),
                tool: None,
                source: ChatMessageSource::Assistant,
                model,
                model_provider: ChatMessageModelProvider::OpenAI,
            },
            OpenAIMessage::Tool {
                content,
                tool_call_id,
                ..
            } => Self {
                role: "tool".to_string(),
                content: Some(ContentTypes::Text(content)),
                tool: Some(Tool {
                    id: tool_call_id,
                    name: None,
                }),
                source: ChatMessageSource::Tool,
                model,
                model_provider: ChatMessageModelProvider::OpenAI,
            },
        }
    }

    pub fn from_anthropic_message(message: AnthropicMessage, model: ChatMessageModel) -> Self {
        match message {
            AnthropicMessage::User { content, .. } => Self {
                role: "user".to_string(),
                content: Some(ContentTypes::List(content)),
                tool: None,
                source: ChatMessageSource::User,
                model,
                model_provider: ChatMessageModelProvider::Anthropic,
            },
            AnthropicMessage::Assistant { content, .. } => {
                if let AnthropicMessageContent::ToolUse { id, name, input: _ } =
                    content.clone().first().unwrap()
                {
                    Self {
                        role: "tool".to_string(),
                        content: Some(ContentTypes::List(content)),
                        tool: Some(Tool {
                            id: id.clone(),
                            name: Some(name.clone()),
                        }),
                        source: ChatMessageSource::Tool,
                        model,
                        model_provider: ChatMessageModelProvider::Anthropic,
                    }
                } else {
                    Self {
                        role: "assistant".to_string(),
                        content: Some(ContentTypes::List(content)),
                        tool: None,
                        source: ChatMessageSource::Assistant,
                        model,
                        model_provider: ChatMessageModelProvider::Anthropic,
                    }
                }
            }
        }
    }
}
