use meridian::llms::{
    anthropic::{messages::AnthropicMessage, AnthropicClient},
    open_ai::{messages::OpenAIMessage, OpenAIClient},
    tools::{Tool, ToolChoice, ToolDefinition, ToolExecution, Toolkit, TypedTool},
    LLMToolUsage,
};
use schemars::{schema::RootSchema, JsonSchema};
use serde::{Deserialize, Serialize};
use serde_json::Value;

use anyhow::Result;

use crate::common::utils::{can_run_anthropic_test, can_run_openai_test, setup_test_env};

#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct CalculatorArgs {
    operation: String,
    numbers: Vec<f64>,
}

#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct CalculatorResult {
    result: f64,
}

struct Calculator;

impl TypedTool for Calculator {
    type Args = CalculatorArgs;
    type Output = CalculatorResult;

    fn name(&self) -> &str {
        "calculator"
    }

    fn description(&self) -> &str {
        "Performs basic arithmetic operations"
    }

    fn execute(&self, args: Self::Args) -> Result<Self::Output> {
        let result = match args.operation.as_str() {
            "add" => args.numbers.iter().sum(),
            "multiply" => args.numbers.iter().product(),
            _ => return Err(anyhow::anyhow!("Unsupported operation")),
        };
        Ok(CalculatorResult { result })
    }
}

#[test]
fn test_anthropic_tool_use() -> Result<()> {
    setup_test_env()?;
    if !can_run_anthropic_test() {
        return Ok(());
    }

    let client = AnthropicClient::new();
    let calculator = Calculator;

    let messages = vec![AnthropicMessage::user(
        "What is 2 + 2? Use the calculator tool.",
    )];

    let mut tool_kit = Toolkit::new();
    tool_kit.add(calculator);

    let response = client.get_chat_with_tools(messages, &tool_kit, &ToolChoice::Any)?;

    assert!(!response.is_empty());
    println!("Tool response: {:?}", response);
    Ok(())
}

#[test]
fn test_anthropic_direct_tool_execution() -> Result<()> {
    setup_test_env()?;
    if !can_run_anthropic_test() {
        return Ok(());
    }

    let client = AnthropicClient::new();
    let calculator = Calculator;

    let messages = vec![AnthropicMessage::user(
        "What is 2 + 2? Use the calculator tool.",
    )];

    let response = client.do_work_with_tool(messages, &calculator)?;

    assert!(!response.is_empty());
    println!("Direct tool response: {:?}", response);
    Ok(())
}

#[test]
fn test_open_ai_tool_use() -> Result<()> {
    setup_test_env()?;
    if !can_run_openai_test() {
        return Ok(());
    }

    let client = OpenAIClient::new();
    let calculator = Calculator;

    let messages = vec![OpenAIMessage::user(
        "What is 2 + 2? Use the calculator tool.",
    )];

    let mut tool_kit = Toolkit::new();
    tool_kit.add(calculator);

    let response = client.get_chat_with_tools(messages, &tool_kit, &ToolChoice::Any)?;

    assert!(!response.is_empty());
    println!("Tool response: {:?}", response);
    Ok(())
}

#[test]
fn test_openai_direct_tool_use() -> Result<()> {
    setup_test_env()?;
    if !can_run_openai_test() {
        return Ok(());
    }

    let client = OpenAIClient::new();
    let calculator = Calculator;

    let messages = vec![OpenAIMessage::user(
        "What is 2 + 2? Use the calculator tool.",
    )];

    let response = client.do_work_with_tool(messages, &calculator)?;

    assert!(!response.is_empty());
    println!("Tool response: {:?}", response);
    Ok(())
}