//! Tiny demo of the LlmClient streaming abstraction.
//!
//! Streams the assistant response token-by-token to stdout. Demonstrates
//! the OpenAI-compatible client; works with OpenAI, Z.AI, vLLM, llama.cpp,
//! or any Chat Completions endpoint.
//!
//! Usage (Z.AI defaults):
//! ```bash
//! export OPENAI_API_KEY=$Z_AI_API_KEY
//! cargo run -p rtx-csm --release --example llm_chat -- \
//!     --base "https://api.z.ai/api/coding/paas/v4" \
//!     --model glm-4.6 \
//!     --prompt "In one short sentence, why is rust good for ML?"
//! ```

use anyhow::Result;
use clap::Parser;
use futures_util::StreamExt;
use rtx_csm::llm_client::{ChatMessage, GenConfig, LlmClient, OpenAiCompatibleClient};
use std::io::Write;

#[derive(Debug, Parser)]
#[command(name = "llm_chat")]
struct Cli {
    /// Base URL of the OpenAI-compatible endpoint. Default OpenAI public.
    #[arg(long, default_value = "https://api.openai.com/v1")]
    base: String,
    /// Model name.
    #[arg(long, default_value = "gpt-4o-mini")]
    model: String,
    /// API key. Reads OPENAI_API_KEY from env if not set.
    #[arg(long)]
    api_key: Option<String>,
    /// User prompt.
    #[arg(long, default_value = "Say hello in one short sentence.")]
    prompt: String,
    /// Optional system prompt.
    #[arg(long, default_value = "You are a concise assistant.")]
    system: String,
    #[arg(long, default_value_t = 0.7)]
    temperature: f32,
    #[arg(long, default_value_t = 256)]
    max_tokens: u32,
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    let api_key = cli
        .api_key
        .or_else(|| std::env::var("OPENAI_API_KEY").ok())
        .ok_or_else(|| anyhow::anyhow!("set --api-key or OPENAI_API_KEY"))?;
    let client = OpenAiCompatibleClient::new(&cli.base, api_key, &cli.model);
    let messages = vec![
        ChatMessage::system(&cli.system),
        ChatMessage::user(&cli.prompt),
    ];
    let config = GenConfig {
        max_tokens: Some(cli.max_tokens),
        temperature: cli.temperature,
        ..GenConfig::default()
    };
    println!("== {} via {} ==", cli.model, cli.base);
    print!("user: {}\nassistant: ", cli.prompt);
    std::io::stdout().flush().ok();

    let t = std::time::Instant::now();
    let mut first_token_ms: Option<u128> = None;
    let mut stream = client.generate_stream(messages, config).await?;
    let mut total_chars = 0;
    while let Some(chunk) = stream.next().await {
        let chunk = chunk?;
        if first_token_ms.is_none() {
            first_token_ms = Some(t.elapsed().as_millis());
        }
        total_chars += chunk.len();
        print!("{chunk}");
        std::io::stdout().flush().ok();
    }
    println!();
    println!(
        "[ttft={}ms total={}ms chars={}]",
        first_token_ms.unwrap_or(0),
        t.elapsed().as_millis(),
        total_chars
    );
    Ok(())
}