use serde::{Deserialize, Serialize};

use crate::types::{
    MessageParam, Model, SystemPrompt, TextBlock, ThinkingConfig, ToolChoice, ToolUnionParam,
};

/// Parameters for counting tokens in messages.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageCountTokensParams {
    /// Input messages.
    ///
    /// Our models are trained to operate on alternating `user` and `assistant`
    /// conversational turns. When creating a new `Message`, you specify the prior
    /// conversational turns with the `messages` parameter, and the model then generates
    /// the next `Message` in the conversation. Consecutive `user` or `assistant` turns
    /// in your request will be combined into a single turn.
    ///
    /// Each input message must be an object with a `role` and `content`. You can
    /// specify a single `user`-role message, or you can include multiple `user` and
    /// `assistant` messages.
    ///
    /// If the final message uses the `assistant` role, the response content will
    /// continue immediately from the content in that message. This can be used to
    /// constrain part of the model's response.
    pub messages: Vec<MessageParam>,

    /// The model that will count tokens.
    ///
    /// See [models](https://docs.anthropic.com/en/docs/models-overview) for additional
    /// details and options.
    pub model: Model,

    /// System prompt.
    ///
    /// A system prompt is a way of providing context and instructions to Claude, such
    /// as specifying a particular goal or role. See the
    /// [guide to system prompts](https://docs.anthropic.com/en/docs/system-prompts).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub system: Option<SystemPrompt>,

    /// Configuration for enabling Claude's extended thinking.
    ///
    /// When enabled, responses include `thinking` content blocks showing Claude's
    /// thinking process before the final answer. Requires a minimum budget of 1,024
    /// tokens and counts towards your `max_tokens` limit.
    ///
    /// See
    /// [extended thinking](https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking)
    /// for details.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub thinking: Option<ThinkingConfig>,

    /// How the model should use the provided tools.
    ///
    /// The model can use a specific tool, any available tool, decide by itself, or not
    /// use tools at all.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<ToolChoice>,

    /// Definitions of tools that the model may use.
    ///
    /// If you include `tools` in your API request, the model may return `tool_use`
    /// content blocks that represent the model's use of those tools. You can then run
    /// those tools using the tool input generated by the model and then optionally
    /// return results back to the model using `tool_result` content blocks.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<ToolUnionParam>>,
}

impl MessageCountTokensParams {
    /// Create a new `MessageCountTokensParams` with the required fields.
    pub fn new(messages: Vec<MessageParam>, model: Model) -> Self {
        Self { messages, model, system: None, thinking: None, tool_choice: None, tools: None }
    }

    /// Add a system prompt as a string.
    pub fn with_system_string(mut self, system: String) -> Self {
        self.system = Some(SystemPrompt::from_string(system));
        self
    }

    /// Add a system prompt as text blocks.
    pub fn with_system_blocks(mut self, blocks: Vec<TextBlock>) -> Self {
        self.system = Some(SystemPrompt::from_blocks(blocks));
        self
    }

    /// Add a system prompt.
    pub fn with_system(mut self, system: SystemPrompt) -> Self {
        self.system = Some(system);
        self
    }

    /// Add thinking configuration.
    pub fn with_thinking(mut self, thinking: ThinkingConfig) -> Self {
        self.thinking = Some(thinking);
        self
    }

    /// Add tool choice.
    pub fn with_tool_choice(mut self, tool_choice: ToolChoice) -> Self {
        self.tool_choice = Some(tool_choice);
        self
    }

    /// Add tools.
    pub fn with_tools(mut self, tools: Vec<ToolUnionParam>) -> Self {
        self.tools = Some(tools);
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{KnownModel, MessageRole};
    use serde_json::{json, to_value};

    #[test]
    fn message_count_tokens_params_serialization() {
        let message = MessageParam::new_with_string("Hello, Claude".to_string(), MessageRole::User);

        let params =
            MessageCountTokensParams::new(vec![message], Model::Known(KnownModel::ClaudeSonnet46));

        let json = to_value(&params).unwrap();
        assert_eq!(
            json,
            json!({
                "messages": [
                    {
                        "role": "user",
                        "content": "Hello, Claude"
                    }
                ],
                "model": "claude-sonnet-4-6"
            })
        );
    }

    #[test]
    fn message_count_tokens_params_with_system() {
        let message = MessageParam::new_with_string("Hello, Claude".to_string(), MessageRole::User);

        let params =
            MessageCountTokensParams::new(vec![message], Model::Known(KnownModel::ClaudeSonnet46))
                .with_system_string("You are a helpful assistant.".to_string());

        let json = to_value(&params).unwrap();
        assert_eq!(
            json,
            json!({
                "messages": [
                    {
                        "role": "user",
                        "content": "Hello, Claude"
                    }
                ],
                "model": "claude-sonnet-4-6",
                "system": "You are a helpful assistant."
            })
        );
    }

    #[test]
    fn message_count_tokens_params_with_tools() {
        let message = MessageParam::new_with_string("Hello, Claude".to_string(), MessageRole::User);

        let tool = ToolUnionParam::new_custom_tool(
            "search".to_string(),
            json!({
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query"
                    }
                }
            }),
        );

        let params =
            MessageCountTokensParams::new(vec![message], Model::Known(KnownModel::ClaudeSonnet46))
                .with_tools(vec![tool]);

        let json = to_value(&params).unwrap();
        assert_eq!(
            json,
            json!({
                "messages": [
                    {
                        "role": "user",
                        "content": "Hello, Claude"
                    }
                ],
                "model": "claude-sonnet-4-6",
                "tools": [
                    {
                        "type": "custom",
                        "input_schema": {
                            "type": "object",
                            "properties": {
                                "query": {
                                    "type": "string",
                                    "description": "The search query"
                                }
                            }
                        },
                        "name": "search"
                    }
                ]
            })
        );
    }
}
