use serde::{Deserialize, Serialize};

use super::AnthropicModel;
use crate::llms::messages::{AbstractMessage, Role};
use anyhow::Result;

#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum AnthropicMessageContent {
    Text {
        text: String,
    },
    String,
    #[serde(rename = "tool_use")]
    ToolUse {
        id: String,
        name: String,
        input: serde_json::Value,
    },
}

#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(tag = "role", rename_all = "lowercase")]
pub enum AnthropicMessage {
    User {
        content: Vec<AnthropicMessageContent>,
        #[serde(skip_serializing_if = "Option::is_none")]
        tool_use_id: Option<String>,
    },
    Assistant {
        content: Vec<AnthropicMessageContent>,
    },
}

impl AnthropicMessage {
    pub fn new(content: Vec<AnthropicMessageContent>, role: Role) -> Self {
        match role {
            Role::User => Self::User {
                content,
                tool_use_id: None,
            },
            Role::Assistant => Self::Assistant { content },
            _ => panic!("Invalid role for Anthropic message"),
        }
    }

    pub fn with_tool_use_id(&self, id: String) -> Self {
        match self {
            AnthropicMessage::User {
                content,
                tool_use_id,
                ..
            } => AnthropicMessage::User {
                content: content.clone(),
                tool_use_id: Some(id),
            },
            _ => self.clone(),
        }
    }

    pub fn user(content: &str) -> Self {
        Self::new(
            vec![AnthropicMessageContent::Text {
                text: content.to_string(),
            }],
            Role::User,
        )
    }

    pub fn assistant(content: &str) -> Self {
        Self::new(
            vec![AnthropicMessageContent::Text {
                text: content.to_string(),
            }],
            Role::Assistant,
        )
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct NextMessageRequest {
    pub model: AnthropicModel,
    pub messages: Vec<AnthropicMessage>,
    pub max_tokens: u64,
}

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

    fn get_content(&self) -> Result<String> {
        match self {
            AnthropicMessage::User { content, .. } => match content.first() {
                Some(AnthropicMessageContent::Text { text, .. }) => Ok(text.clone()),
                _ => Err(anyhow::anyhow!("No text content")),
            },
            AnthropicMessage::Assistant { content, .. } => match content.first() {
                Some(AnthropicMessageContent::Text { text, .. }) => Ok(text.clone()),
                _ => Err(anyhow::anyhow!("No text content")),
            },
        }
    }
}
