use rig_core::client::{CompletionClient, ProviderClient};
use rig_core::completion::Prompt;
use rig_core::providers;
use rig_core::tool::{Tool, ToolError};
use rig_derive::rig_tool;
use std::time::Duration;

/// A tool that simulates an async operation
#[rig_tool]
async fn async_operation(
    /// Input value to process
    input: String,
    /// Delay in milliseconds before returning result
    delay_ms: u64,
) -> Result<String, ToolError> {
    tokio::time::sleep(Duration::from_millis(delay_ms)).await;

    Ok(format!(
        "Processed after {}ms: {}",
        delay_ms,
        input.to_uppercase()
    ))
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    tracing_subscriber::fmt().pretty().init();

    let async_agent = providers::openai::Client::from_env()?
        .agent(providers::openai::GPT_4O)
        .preamble("You are an agent with tools access, always use the tools")
        .max_tokens(1024)
        .tool(AsyncOperation)
        .build();

    println!("Tool definition:");
    println!(
        "ASYNCOPERATION: {}",
        serde_json::to_string_pretty(&AsyncOperation.definition(String::default()).await)?
    );

    for prompt in [
        "What tools do you have?",
        "Process the text 'hello world' with a delay of 1000ms",
        "Process the text 'async operation' with a delay of 500ms",
        "Process the text 'concurrent calls' with a delay of 200ms",
        "Process the text 'error handling' with a delay of 'not a number'",
    ] {
        println!("User: {prompt}");
        println!("Agent: {}", async_agent.prompt(prompt).await?);
    }

    Ok(())
}
