pub mod messages;
use anyhow::Result;
use log::{debug, info, warn};
use messages::{AnthropicMessage, AnthropicMessageContent, NextMessageRequest};
use reqwest::blocking::Client;
use schemars::{schema_for, JsonSchema};
use serde::{Deserialize, Serialize};

use super::{
    messages::{Message, Role},
    LLMProvider, LLMToolUsage, MultiModelLLMProvider, StructuredLLMProvider, Tool, ToolChoice,
    ToolDefinition, ToolExecution, Toolkit,
};

#[derive(Debug, Serialize, Deserialize, Clone, Copy)]
pub enum AnthropicModel {
    #[serde(rename = "claude-3-5-sonnet-20241022")]
    Claude35Sonnet,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
struct Usage {
    input_tokens: u64,
    output_tokens: u64,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct NextMessageResponse {
    id: String,
    #[serde(rename = "type")]
    r#type: String,
    role: String,
    model: AnthropicModel,
    usage: Usage,
    content: Vec<AnthropicMessageContent>,
    stop_reason: String,
    stop_sequence: Option<String>,
}

pub struct AnthropicClient {
    api_key: String,
    version: String,
    client: Client,
    model: AnthropicModel,
}

impl Default for AnthropicClient {
    fn default() -> Self {
        Self::new()
    }
}

impl AnthropicClient {
    pub fn new() -> Self {
        Self {
            api_key: std::env::var("ANTHROPIC_API_KEY").unwrap(),
            version: "2023-06-01".to_string(),
            client: Client::new(),
            model: AnthropicModel::Claude35Sonnet,
        }
    }
}

impl LLMProvider<AnthropicMessage> for AnthropicClient {
    fn get_completion(&self, messages: Vec<AnthropicMessage>) -> Result<Vec<AnthropicMessage>> {
        debug!(
            "Getting completion from Anthropic with {} messages",
            messages.len()
        );

        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert("x-api-key", self.api_key.parse().unwrap());
        headers.insert("anthropic-version", self.version.parse().unwrap());
        headers.insert("content-type", "application/json".parse().unwrap());

        let request_body = NextMessageRequest {
            model: self.model,
            messages: messages.clone(),
            max_tokens: 8192,
        };

        debug!("Sending request to Anthropic API");
        let result = self
            .client
            .post("https://api.anthropic.com/v1/messages")
            .headers(headers)
            .json(&request_body)
            .send()?;

        if !result.status().is_success() {
            let status = result.status();
            let error_text = result.text()?;
            warn!("Anthropic API error: {} - {}", status, error_text);
            return Err(anyhow::anyhow!("Failed to get completion"));
        }

        let response: NextMessageResponse = result.json()?;
        debug!("Received response from Anthropic API");
        debug!("Last message: {:?}", response.content.last());

        let mut all_messages = messages;
        all_messages.append(&mut vec![response.get_latest_message()?]);
        Ok(all_messages)
    }

    fn stream_completion(
        &self,
        messages: Vec<AnthropicMessage>,
    ) -> Result<Box<dyn Iterator<Item = AnthropicMessage>>> {
        todo!("Implement streaming for the Anthropic client")
    }
}

impl MultiModelLLMProvider<AnthropicModel> for AnthropicClient {
    fn with_model(&self, model: AnthropicModel) -> Self {
        Self {
            api_key: self.api_key.clone(),
            version: self.version.clone(),
            client: self.client.clone(),
            model,
        }
    }

    fn get_model(&self) -> AnthropicModel {
        self.model
    }
}

impl StructuredLLMProvider<AnthropicMessage> for AnthropicClient {
    fn get_structured_response<
        'de,
        DesiredSchema: Serialize + serde::de::DeserializeOwned + JsonSchema,
    >(
        &self,
        messages: Vec<AnthropicMessage>,
    ) -> Result<DesiredSchema> {
        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert("x-api-key", self.api_key.parse().unwrap());
        headers.insert("anthropic-version", self.version.parse().unwrap());
        headers.insert("content-type", "application/json".parse().unwrap());

        let tool = serde_json::json!({
            "name": "process_structured_output",
            "description": "Process the structured output response",
            "input_schema": schema_for!(DesiredSchema).serialize(serde_json::value::Serializer)?,
        });

        let request_body = serde_json::json!({
            "model": AnthropicModel::Claude35Sonnet,
            "messages": messages,
            "max_tokens": 8192,
            "tools": [tool],
            "tool_choice": {
                "type": "tool",
                "name": "process_structured_output"
            }
        });

        let result = self
            .client
            .post("https://api.anthropic.com/v1/messages")
            .headers(headers)
            .json(&request_body)
            .send()?;

        if !result.status().is_success() {
            println!("{}", result.text()?);
            return Err(anyhow::anyhow!("Failed to get structured response"));
        }

        let response: serde_json::Value = result.json()?;


        // Extract tool call response from the content array
        let tool_response = response["content"]
            .as_array()
            .and_then(|content| content.first())
            .and_then(|tool_call| tool_call.get("input"))
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("Failed to get tool response"))?;

        // Deserialize into the desired schema
        Ok(serde_json::from_value(tool_response)?)
    }
}

impl NextMessageResponse {
    pub fn get_latest_message(&self) -> Result<AnthropicMessage> {
        let message_content = self.content.last().unwrap();
        Ok(AnthropicMessage::Assistant {
            content: vec![message_content.clone()],
        })
    }
}

impl LLMToolUsage<AnthropicMessage> for AnthropicClient {
    fn do_work_with_tool(
        &self,
        messages: Vec<AnthropicMessage>,
        tool: &dyn Tool,
    ) -> Result<Vec<AnthropicMessage>> {
        debug!("Executing tool '{}' with Anthropic", tool.name());

        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert("x-api-key", self.api_key.parse().unwrap());
        headers.insert("anthropic-version", self.version.parse().unwrap());
        headers.insert("content-type", "application/json".parse().unwrap());

        let request_body = serde_json::json!({
            "model": self.model,
            "messages": messages,
            "max_tokens": 8192,
            "tools": [{
                "name": tool.name(),
                "description": tool.description(),
                "input_schema": tool.schema(),
            }],
            "tool_choice": {
                "type": "tool",
                "name": tool.name()
            }
        });

        debug!("Sending tool execution request to Anthropic API");
        let result = self
            .client
            .post("https://api.anthropic.com/v1/messages")
            .headers(headers)
            .json(&request_body)
            .send()?;

        if !result.status().is_success() {
            let status = result.status();
            let error_text = result.text()?;
            warn!(
                "Anthropic API error during tool execution: {} - {}",
                status, error_text
            );
            return Err(anyhow::anyhow!("Failed to use tool"));
        }

        let response: NextMessageResponse = result.json()?;
        debug!("Received tool execution response from Anthropic");
        debug!("Last message content: {:?}", response.content.first());

        if let Some(content) = response.content.first() {
            match content {
                AnthropicMessageContent::ToolUse { input, id, .. } => {
                    debug!("Processing tool use response with id: {}", id);
                    let result = tool.execute(input.clone())?;
                    Ok(vec![AnthropicMessage::new(
                        vec![AnthropicMessageContent::Text {
                            text: serde_json::to_string(&result)?,
                        }],
                        Role::Assistant,
                    )
                    .with_tool_use_id(id.clone())])
                }
                _ => {
                    warn!("Unexpected response content type from Anthropic");
                    Err(anyhow::anyhow!("Unexpected response content type"))
                }
            }
        } else {
            warn!("No response content from Anthropic");
            Err(anyhow::anyhow!("No response content"))
        }
    }

    /// Calls the LLM with tools and returns the next message in the conversation
    ///
    /// This method does not run the tool, it just adds the next message to the conversation
    fn get_chat_with_tools(
        &self,
        messages: Vec<AnthropicMessage>,
        tool_kit: &Toolkit,
        force_tool_use: &ToolChoice,
    ) -> Result<Vec<AnthropicMessage>> {
        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert("x-api-key", self.api_key.parse().unwrap());
        headers.insert("anthropic-version", self.version.parse().unwrap());
        headers.insert("content-type", "application/json".parse().unwrap());

        let tool_defs: Vec<serde_json::Value> = tool_kit
            .tools()
            .iter()
            .map(|tool| {
                serde_json::json!({
                    "name": tool.name(),
                    "description": tool.description(),
                    "input_schema": tool.schema()
                })
            })
            .collect();

        let tool_choice = match force_tool_use {
            ToolChoice::Any => serde_json::json!({ "type": "any" }),
            ToolChoice::Specific(name) => serde_json::json!({ "type": "tool", "name": name }),
            ToolChoice::SelfSelect => serde_json::json!({ "type": "auto" }),
        };

        let request_body = serde_json::json!({
            "model": self.model,
            "messages": messages,
            "max_tokens": 8192,
            "tools": tool_defs,
            "tool_choice": tool_choice
        });

        let result = self
            .client
            .post("https://api.anthropic.com/v1/messages")
            .headers(headers)
            .json(&request_body)
            .send()?;

        let response: NextMessageResponse = result.json()?;
        debug!("Last message: {:?}", response.content.last());
        Ok(vec![response.get_latest_message()?])
    }

    /// Gets a completion with tools and then runs the tool
    fn get_work_result(
        &self,
        messages: Vec<AnthropicMessage>,
        tool_kit: &Toolkit,
        tool_choice: &ToolChoice,
    ) -> Result<Vec<AnthropicMessage>> {
        info!("Getting work result with tool choice: {:?}", tool_choice);

        match tool_choice {
            ToolChoice::Specific(name) => {
                debug!("Using specific tool: {}", name);
                self.do_work_with_tool(
                    messages,
                    tool_kit
                        .get(name)
                        .ok_or_else(|| anyhow::anyhow!("Tool not found: {}", name))?,
                )
            }
            ToolChoice::Any => {
                debug!("Getting chat with any tool allowed");
                let response = self.get_chat_with_tools(messages, tool_kit, tool_choice)?;

                if let Some(AnthropicMessage::Assistant { content }) = response.clone().last() {
                    debug!("Assistant message content: {:?}", content);
                    let mut result_messages = response;

                    // Process all tool calls in the content
                    for message_content in content {
                        if let AnthropicMessageContent::ToolUse {
                            input, id, name, ..
                        } = message_content
                        {
                            debug!("Processing tool use - Name: {}, ID: {}", name, id);
                            let tool = tool_kit
                                .get(name)
                                .ok_or_else(|| anyhow::anyhow!("Tool not found: {}", name))?;

                            let result = tool.execute(input.clone())?;
                            debug!("Tool execution result: {:?}", result);

                            result_messages.push(
                                AnthropicMessage::new(
                                    vec![AnthropicMessageContent::Text {
                                        text: serde_json::to_string(&result)?,
                                    }],
                                    Role::Assistant,
                                )
                                .with_tool_use_id(id.clone()),
                            );
                        }
                    }
                    Ok(result_messages)
                } else {
                    Err(anyhow::anyhow!("No tool calls in assistant message"))
                }
            }
            ToolChoice::SelfSelect => {
                debug!("Letting model select tool usage");
                let response = self.get_chat_with_tools(messages, tool_kit, tool_choice)?;

                if let Some(AnthropicMessage::Assistant { content }) = response.clone().last() {
                    debug!("Assistant message content: {:?}", content);
                    let mut result_messages = response;

                    // Process all tool calls in the content
                    for message_content in content {
                        if let AnthropicMessageContent::ToolUse {
                            input, id, name, ..
                        } = message_content
                        {
                            debug!("Processing tool use - Name: {}, ID: {}", name, id);
                            let tool = tool_kit
                                .get(name)
                                .ok_or_else(|| anyhow::anyhow!("Tool not found: {}", name))?;

                            let result = tool.execute(input.clone())?;
                            debug!("Tool execution result: {:?}", result);

                            result_messages.push(
                                AnthropicMessage::new(
                                    vec![AnthropicMessageContent::Text {
                                        text: serde_json::to_string(&result)?,
                                    }],
                                    Role::Assistant,
                                )
                                .with_tool_use_id(id.clone()),
                            );
                        }
                    }
                    Ok(result_messages)
                } else {
                    Ok(response) // For SelfSelect, we return the response even if no tools were used
                }
            }
        }
    }
}

impl Message for AnthropicMessage {
    fn new(content: &str, role: Role) -> Self {
        match role {
            Role::User => Self::User {
                content: vec![AnthropicMessageContent::Text {
                    text: content.to_string(),
                }],
                tool_use_id: None,
            },
            Role::Assistant => Self::Assistant {
                content: vec![AnthropicMessageContent::Text {
                    text: content.to_string(),
                }],
            },
            _ => panic!("Invalid role for Anthropic message"),
        }
    }
}
