//! Provider Comparison Demo
//!
//! Demonstrates all three LLM providers (Claude, OpenAI, Gemini) with the same prompt,
//! showing both blocking and streaming modes.
//!
//! Usage:
//!   # Run all providers
//!   cargo run --example provider_comparison
//!
//!   # Run specific provider
//!   cargo run --example provider_comparison claude
//!   cargo run --example provider_comparison openai
//!   cargo run --example provider_comparison gemini
//!
//! Environment variables required:
//!   ANTHROPIC_API_KEY - for Claude
//!   OPENAI_API_KEY - for OpenAI
//!   GEMINI_API_KEY or GOOGLE_API_KEY - for Gemini

use futures::stream::StreamExt;
use std::io::{self, Write};
use std::time::Instant;
use typedialog_ag_core::llm::{
    ClaudeProvider, GeminiProvider, LlmMessage, LlmProvider, LlmRequest, MessageRole,
    OllamaProvider, OpenAIProvider, StreamChunk,
};

const PROMPT: &str =
    "Write a haiku about artificial intelligence. Return only the haiku, nothing else.";

struct ProviderConfig {
    name: &'static str,
    model: &'static str,
    color: &'static str,
}

const CLAUDE_CONFIG: ProviderConfig = ProviderConfig {
    name: "Claude",
    model: "claude-3-5-haiku-20241022",
    color: "\x1b[35m", // Magenta
};

const OPENAI_CONFIG: ProviderConfig = ProviderConfig {
    name: "OpenAI",
    model: "gpt-4o-mini",
    color: "\x1b[32m", // Green
};

const GEMINI_CONFIG: ProviderConfig = ProviderConfig {
    name: "Gemini",
    model: "gemini-2.0-flash-exp",
    color: "\x1b[34m", // Blue
};

const OLLAMA_CONFIG: ProviderConfig = ProviderConfig {
    name: "Ollama",
    model: "llama2",
    color: "\x1b[33m", // Yellow
};

const RESET: &str = "\x1b[0m";
const BOLD: &str = "\x1b[1m";
const DIM: &str = "\x1b[2m";

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    let provider_filter = args.get(1).map(|s| s.as_str());

    println!("{}{}", BOLD, "═".repeat(70));
    println!("🤖  TypeAgent Multi-Provider Demo");
    println!("{}{}\n", "═".repeat(70), RESET);

    println!("{}Prompt:{} {}\n", BOLD, RESET, PROMPT);

    // Determine which providers to run
    let run_claude = provider_filter.is_none() || provider_filter == Some("claude");
    let run_openai = provider_filter.is_none() || provider_filter == Some("openai");
    let run_gemini = provider_filter.is_none() || provider_filter == Some("gemini");
    let run_ollama = provider_filter.is_none() || provider_filter == Some("ollama");

    // Run blocking mode comparison
    println!("{}{}", BOLD, "━".repeat(70));
    println!("📋  BLOCKING MODE (complete() method)");
    println!("{}{}\n", "━".repeat(70), RESET);

    if run_claude {
        run_blocking_demo(&CLAUDE_CONFIG, ClaudeProvider::new()).await;
    }

    if run_openai {
        run_blocking_demo(&OPENAI_CONFIG, OpenAIProvider::new()).await;
    }

    if run_gemini {
        run_blocking_demo(&GEMINI_CONFIG, GeminiProvider::new()).await;
    }

    if run_ollama {
        run_blocking_demo(&OLLAMA_CONFIG, OllamaProvider::new()).await;
    }

    // Run streaming mode comparison
    println!("\n{}{}", BOLD, "━".repeat(70));
    println!("⚡  STREAMING MODE (stream() method)");
    println!("{}{}\n", "━".repeat(70), RESET);

    if run_claude {
        run_streaming_demo(&CLAUDE_CONFIG, ClaudeProvider::new()).await;
    }

    if run_openai {
        run_streaming_demo(&OPENAI_CONFIG, OpenAIProvider::new()).await;
    }

    if run_gemini {
        run_streaming_demo(&GEMINI_CONFIG, GeminiProvider::new()).await;
    }

    if run_ollama {
        run_streaming_demo(&OLLAMA_CONFIG, OllamaProvider::new()).await;
    }

    // Summary
    println!("\n{}{}", BOLD, "═".repeat(70));
    println!("📊  SUMMARY");
    println!("{}{}\n", "═".repeat(70), RESET);

    println!("{}All four providers support:{}", BOLD, RESET);
    println!("  ✅ Blocking requests (complete())");
    println!("  ✅ Streaming responses (stream())");
    println!("  ✅ Token usage tracking");
    println!("  ✅ Error handling");
    println!("  ✅ System message support");
    println!("\n{}Provider-specific features:{}", BOLD, RESET);
    println!("  🟣 Claude:  SSE streaming, usage in stream, cloud");
    println!("  🟢 OpenAI:  SSE streaming, no usage in stream, cloud");
    println!("  🔵 Gemini:  JSON streaming, usage in stream, cloud, 'model' role");
    println!("  🟡 Ollama:  JSON streaming, usage in stream, local models");

    Ok(())
}

async fn run_blocking_demo<P: LlmProvider>(
    config: &ProviderConfig,
    provider_result: Result<P, typedialog_ag_core::error::Error>,
) {
    println!(
        "{}{}{} ({}){}",
        config.color, BOLD, config.name, config.model, RESET
    );
    println!(
        "{}────────────────────────────────────────────────────────────────────{}",
        DIM, RESET
    );

    let provider = match provider_result {
        Ok(p) => p,
        Err(e) => {
            println!("{}❌ Error: {}{}\n", config.color, e, RESET);
            return;
        }
    };

    let request = LlmRequest {
        model: config.model.to_string(),
        messages: vec![LlmMessage {
            role: MessageRole::User,
            content: PROMPT.to_string(),
        }],
        max_tokens: Some(100),
        temperature: Some(0.7),
        system: Some("You are a creative poet.".to_string()),
    };

    let start = Instant::now();

    match provider.complete(request).await {
        Ok(response) => {
            let duration = start.elapsed();

            println!("{}Response:{}", BOLD, RESET);
            println!("{}{}{}", config.color, response.content, RESET);

            println!("\n{}Metadata:{}", DIM, RESET);
            println!("  Duration: {}ms", duration.as_millis());
            if let Some(usage) = response.usage {
                println!(
                    "  Tokens: {} (in: {}, out: {})",
                    usage.total_tokens, usage.input_tokens, usage.output_tokens
                );
            }
            println!();
        }
        Err(e) => {
            println!("{}❌ Error: {}{}\n", config.color, e, RESET);
        }
    }
}

async fn run_streaming_demo<P: LlmProvider>(
    config: &ProviderConfig,
    provider_result: Result<P, typedialog_ag_core::error::Error>,
) {
    println!(
        "{}{}{} ({}){}",
        config.color, BOLD, config.name, config.model, RESET
    );
    println!(
        "{}────────────────────────────────────────────────────────────────────{}",
        DIM, RESET
    );

    let provider = match provider_result {
        Ok(p) => p,
        Err(e) => {
            println!("{}❌ Error: {}{}\n", config.color, e, RESET);
            return;
        }
    };

    let request = LlmRequest {
        model: config.model.to_string(),
        messages: vec![LlmMessage {
            role: MessageRole::User,
            content: PROMPT.to_string(),
        }],
        max_tokens: Some(100),
        temperature: Some(0.7),
        system: Some("You are a creative poet.".to_string()),
    };

    let start = Instant::now();

    match provider.stream(request).await {
        Ok(mut stream) => {
            print!("{}Response: {}{}", BOLD, RESET, config.color);
            io::stdout().flush().unwrap();

            let mut full_output = String::new();
            let mut tokens = None;

            while let Some(chunk_result) = stream.next().await {
                match chunk_result {
                    Ok(StreamChunk::Content(text)) => {
                        print!("{}", text);
                        io::stdout().flush().unwrap();
                        full_output.push_str(&text);
                    }
                    Ok(StreamChunk::Done(metadata)) => {
                        tokens = metadata.usage.map(|u| u.total_tokens);
                    }
                    Ok(StreamChunk::Error(err)) => {
                        println!("\n{}❌ Stream error: {}{}", config.color, err, RESET);
                    }
                    Err(e) => {
                        println!("\n{}❌ Error: {}{}", config.color, e, RESET);
                    }
                }
            }

            let duration = start.elapsed();

            println!("{}", RESET);
            println!("\n{}Metadata:{}", DIM, RESET);
            println!("  Duration: {}ms", duration.as_millis());
            if let Some(t) = tokens {
                println!("  Tokens: {}", t);
            }
            println!("  Characters streamed: {}", full_output.len());
            println!();
        }
        Err(e) => {
            println!("{}❌ Error: {}{}\n", config.color, e, RESET);
        }
    }
}
