use axllm::{ax, AxResult, OpenAICompatibleClient};
use serde_json::json;
use std::env;


fn openai_client() -> AxResult<OpenAICompatibleClient> {
    let api_key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY")).map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.6-luna".to_string());
    Ok(OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0})))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;
    let mut program = ax("question:string -> answer:string")?;
    let output = program.forward_with_options(
        &mut client,
        json!({"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."}),
        json!({"promptCacheKey": "ax-openai-example", "contextCache": {}}),
    )?;
    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}

use axllm::{get_supported_ai_models, AxResult};
use serde_json::{json, Value};

fn provider<'a>(catalog: &'a [Value], name: &str) -> &'a Value {
    catalog
        .iter()
        .find(|entry| entry["name"] == name)
        .unwrap_or_else(|| panic!("missing provider {name}"))
}

fn main() -> AxResult<()> {
    let catalog = get_supported_ai_models()?;
    let azure = provider(&catalog, "azure-openai");
    let openrouter = provider(&catalog, "openrouter");

    assert_eq!(azure["isDynamic"], true);
    assert_eq!(azure["models"], json!([]));
    assert!(azure["capabilities"]["thinkingLevels"]
        .as_array()
        .is_some_and(|levels| levels.iter().any(|level| level == "high")));
    assert!(azure["capabilities"]["serviceTiers"]
        .as_array()
        .is_some_and(|tiers| tiers.iter().any(|tier| tier == "priority")));
    assert!(openrouter["capabilities"]["serviceTiers"]
        .as_array()
        .is_some_and(|tiers| tiers.iter().any(|tier| tier == "flex")));

    println!(
        "{} providers; Azure and OpenRouter named profiles are available",
        catalog.len()
    );
    Ok(())
}

use axllm::{ax, AxResult, OpenAICompatibleClient};
use serde_json::json;
use std::env;


fn openai_client() -> AxResult<OpenAICompatibleClient> {
    let api_key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY")).map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".to_string());
    Ok(OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0})))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;
    let mut program = ax("ticket:string -> priority:class \"high, normal, low\", summary:string, labels:string[]")?;
    let output = program.forward(&mut client, json!({"ticket": "Checkout has failed for enterprise customers since 09:00. Support wants a concise summary and tags."}))?;
    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}

use axllm::{ai, AxAIClient, AxError, AxResult};
use serde_json::json;
use std::env;

fn required(name: &str) -> AxResult<String> {
    env::var(name).map_err(|_| AxError::runtime(format!("Set {name} to run this example.")))
}

fn main() -> AxResult<()> {
    let model = env::var("AX_VERTEX_MODEL").unwrap_or_else(|_| "gemini-3.5-flash".to_string());
    let mut client = ai("google-gemini", json!({
        "api_key": required("GOOGLE_VERTEX_ACCESS_TOKEN")?,
        "project_id": required("GOOGLE_PROJECT_ID")?,
        "region": required("GOOGLE_REGION")?,
        "model": model,
    }))?;
    let out = client.chat(json!({
        "chat_prompt": [{"role": "user", "content": "Reply with the word ready."}]
    }))?;
    println!("{}", serde_json::to_string_pretty(&out)?);
    Ok(())
}

use axllm::{ax, f, AxResult, FieldType, OpenAICompatibleClient};
use serde_json::json;
use std::env;

fn openai_client() -> AxResult<OpenAICompatibleClient> {
    let api_key = env::var("OPENAI_API_KEY")
        .or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| {
            axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")
        })?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".to_string());
    Ok(OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0})))
}

fn main() -> AxResult<()> {
    let mut request_type = FieldType::string();
    request_type.min_length = Some(10.0);
    request_type.max_length = Some(500.0);
    request_type.description = Some("Booking request".to_string());
    let mut email_type = FieldType::string();
    email_type.format = Some("email".to_string());
    email_type.description = Some("Contact email".to_string());
    let mut party_type = FieldType::number();
    party_type.minimum = Some(1.0);
    party_type.maximum = Some(12.0);
    party_type.description = Some("Guests".to_string());
    let mut code_type = FieldType::string();
    code_type.pattern = Some(r"^[A-Z]{3}-\d{4}$".to_string());
    code_type.pattern_description = Some("Must look like ABC-1234".to_string());

    let signature = f()
        .input("requestText", request_type)
        .input("contactEmail", email_type)
        .output("partySize", party_type)
        .output("bookingCode", code_type)
        .build();
    let mut program = ax("requestText:string -> partySize:number, bookingCode:string")?;
    program.signature = signature;
    let mut client = openai_client()?;
    let output = program.forward(
        &mut client,
        json!({
            "requestText": "Book dinner for four people under the name Ada Lovelace.",
            "contactEmail": "ada@example.com"
        }),
    )?;
    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}

use axllm::{
    set_usage_observer, AxAIClient, AxError, AxResult, AxUsageEvent, OpenAICompatibleClient,
};
use serde_json::json;
use std::env;
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};

fn main() -> AxResult<()> {
    let api_key = env::var("OPENAI_API_KEY")
        .or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".to_string());
    let events = Arc::new(Mutex::new(Vec::<AxUsageEvent>::new()));
    let captured = Arc::clone(&events);
    set_usage_observer(Some(Arc::new(move |event| {
        captured.lock().unwrap().push(event);
    })));

    let mut client = OpenAICompatibleClient::new(api_key, model).with_options(json!({
        "usageContext": {
            "tenantId": "tenant-42",
            "feature": "support-chat",
            "attributes": {"environment": "example"}
        }
    }));
    let request_id = format!(
        "request-{}",
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_err(|error| AxError::runtime(error.to_string()))?
            .as_nanos()
    );
    client.chat_with_options(
        json!({
            "chat_prompt": [
                {"role": "user", "content": "Reply with one short greeting."}
            ]
        }),
        json!({
            "usageContext": {"userId": "user-7", "requestId": request_id}
        }),
    )?;
    set_usage_observer(None);
    println!("{}", serde_json::to_string_pretty(&*events.lock().unwrap())?);
    Ok(())
}

use axllm::{ai, AxAIClient, AxResult};
use serde_json::json;
use std::{env, time::Instant};

fn main() -> AxResult<()> {
    let api_key = env::var("OPENAI_API_KEY")
        .or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| {
            axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")
        })?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.6-luna".to_string());
    let mut client = ai(
        "openai",
        json!({"api_key": api_key, "model": model}),
    )?;
    let started = Instant::now();
    for event in client.stream_iter(json!({
        "chat_prompt": [{"role": "user", "content": "Reply with exactly: streaming works"}],
        "model_config": {"temperature": 1}
    }))? {
        let event = event?;
        if let Some(content) = event["results"][0]["content"]
            .as_str()
            .filter(|value| !value.is_empty())
        {
            print!("[{} ms] {content}", started.elapsed().as_millis());
        }
    }
    println!();
    Ok(())
}

use axllm::{ai, AxAIClient, AxError, AxResult};
use serde_json::json;
use std::env;

fn api_key() -> AxResult<String> {
    env::var("GOOGLE_API_KEY")
        .or_else(|_| env::var("GOOGLE_APIKEY"))
        .map_err(|_| AxError::runtime("Set GOOGLE_API_KEY or GOOGLE_APIKEY to run this example."))
}

fn main() -> AxResult<()> {
    let model = env::var("AX_GEMINI_MODEL").unwrap_or_else(|_| "gemini-3.7-flash".to_string());
    let mut client = ai(
        "google-gemini",
        json!({
            "api_key": api_key()?,
            "model": model,
        }),
    )?;
    let out = client.chat_with_options(
        json!({
            "chat_prompt": [{
                "role": "user",
                "content": "Explain in one sentence why batch evaluations save time."
            }]
        }),
        json!({"service_tier": "flex"}),
    )?;
    println!("{}", serde_json::to_string_pretty(&out)?);
    Ok(())
}

use axllm::{ax, AxResult, OpenAICompatibleClient};
use serde_json::json;
use std::env;


fn openai_client() -> AxResult<OpenAICompatibleClient> {
    let api_key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY")).map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".to_string());
    Ok(OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0})))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;
    let mut program = ax("context:string, question:string -> answer:string, citations:string[]")?;
    let output = program.forward(&mut client, json!({"context": "Ax uses signatures, ai(), ax(), agent(), flow(), and optimize().", "question": "How should a new developer think about Ax?"}))?;
    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}

use std::{env, sync::{Arc, Mutex}};

use axllm::{
    AxAIClient, AxBalancer, AxBalancerAdaptiveStrategy, AxBalancerOptions,
    AxInMemoryBalancerStatsStore, AxResult, OpenAICompatibleClient,
};
use serde_json::json;

fn main() -> AxResult<()> {
    let key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".into());
    let clients: Vec<Box<dyn AxAIClient>> = vec![
        Box::new(OpenAICompatibleClient::new(&key, &model)),
        Box::new(OpenAICompatibleClient::new(&key, &model)),
    ];

    let store = Arc::new(AxInMemoryBalancerStatsStore::new());
    let route_keys = ["openai-primary".to_string(), "openai-backup".to_string()];
    let events = Arc::new(Mutex::new(Vec::new()));
    let event_sink = events.clone();
    let strategy = AxBalancerAdaptiveStrategy::new(6_000.0, 0.02)
        .with_expected_tokens(1_200, 300)
        .with_namespace("support-summary-v1")
        .with_store(store)
        .with_route_key(Arc::new(move |_service, index| route_keys[index].clone()))
        .with_slice(Arc::new(|context| if context["options"]["stream"] == true { "streaming".into() } else { "interactive".into() }))
        .on_routing_event(Arc::new(move |event| event_sink.lock().unwrap().push(event["type"].clone())));
    let mut balancer = AxBalancer::from_clients(clients, AxBalancerOptions { strategy: Some(strategy), ..AxBalancerOptions::default() })?;
    let response = balancer.chat(json!({"model": model, "chat_prompt": [{"role": "user", "content": "Summarize why shared routing state matters."}]}))?;
    println!("{}", serde_json::to_string_pretty(&response)?);
    println!("{:?}", events.lock().unwrap());
    Ok(())
}

use axllm::runtime::quickjs::QuickJsCodeRuntime;
use axllm::{
    agent_with_options, ax, flow, set_meter, set_rate_limiter, set_tracer, AxCounter,
    AxError, AxGauge, AxHistogram, AxMeter, AxMetricInstrumentOptions, AxRateLimitInfo,
    AxRateLimiter, AxResult, AxRuntimeHooks, AxSpan, AxSpanStart, AxTracer,
    OpenAICompatibleClient,
};
use serde_json::{json, Value};
use std::collections::BTreeMap;
use std::env;
use std::fmt;
use std::sync::Arc;

struct LogSpan(String);
impl fmt::Debug for LogSpan { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_tuple("LogSpan").field(&self.0).finish() } }
impl AxSpan for LogSpan {
    fn add_event(&self, name: &str, _: &BTreeMap<String, Value>) { println!("[span:event] {} {}", self.0, name); }
    fn record_exception(&self, error: &AxError) { println!("[span:error] {} {}", self.0, error); }
    fn end(&self) { println!("[span:end] {}", self.0); }
}

struct LogTracer;
impl AxTracer for LogTracer {
    fn start_span(&self, start: AxSpanStart) -> Option<Arc<dyn AxSpan>> {
        println!("[span:start] {}", start.name);
        Some(Arc::new(LogSpan(start.name)))
    }
}

struct LogInstrument(String);
impl AxCounter for LogInstrument { fn add(&self, value: f64, _: &BTreeMap<String, Value>) { println!("[metric] {} += {}", self.0, value); } }
impl AxHistogram for LogInstrument { fn record(&self, value: f64, _: &BTreeMap<String, Value>) { println!("[metric] {} = {}", self.0, value); } }
impl AxGauge for LogInstrument { fn record(&self, value: f64, _: &BTreeMap<String, Value>) { println!("[metric] {} = {}", self.0, value); } }

struct LogMeter;
impl AxMeter for LogMeter {
    fn create_counter(&self, name: &str, _: &AxMetricInstrumentOptions) -> Option<Arc<dyn AxCounter>> { Some(Arc::new(LogInstrument(name.into()))) }
    fn create_histogram(&self, name: &str, _: &AxMetricInstrumentOptions) -> Option<Arc<dyn AxHistogram>> { Some(Arc::new(LogInstrument(name.into()))) }
    fn create_gauge(&self, name: &str, _: &AxMetricInstrumentOptions) -> Option<Arc<dyn AxGauge>> { Some(Arc::new(LogInstrument(name.into()))) }
}

fn limiter(label: &'static str) -> Arc<dyn AxRateLimiter> {
    Arc::new(move |next: &mut dyn FnMut() -> AxResult<Value>, info: &AxRateLimitInfo| {
        println!("[limit:{label}] {} {}/{} stream={}", info.operation, info.provider, info.model, info.streaming);
        next()
    })
}

fn main() -> AxResult<()> {
    let api_key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".into());
    let mut client = OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0}));
    let tracer: Arc<dyn AxTracer> = Arc::new(LogTracer);
    let meter: Arc<dyn AxMeter> = Arc::new(LogMeter);
    let override_hooks = AxRuntimeHooks { rate_limiter: Some(limiter("forward")), tracer: Some(Arc::clone(&tracer)), meter: Some(Arc::clone(&meter)) };

    set_rate_limiter(Some(limiter("global")));
    set_tracer(Some(Arc::clone(&tracer)));
    set_meter(Some(Arc::clone(&meter)));
    let result = (|| -> AxResult<()> {
        println!("{}", ax("topic:string -> summary:string")?.forward(&mut client, json!({"topic": "portable Ax runtime hooks"}))?);

        let mut helper = agent_with_options("question:string -> answer:string", json!({}))?
            .with_runtime(Box::new(QuickJsCodeRuntime::new()))?;
        println!("{}", helper.forward_with_hooks(&mut client, json!({"question": "What does a rate limiter wrap?"}), json!({"max_actor_steps": 12}), override_hooks.clone())?);

        let mut workflow = flow("examples.runtimeHooks")
            .execute("outline", ax("topic:string -> outline:string")?)
            .execute("polish", ax("outline:string -> answer:string")?)
            .returns(json!({"answer": "polish"}));
        println!("{}", workflow.forward_with_hooks(&mut client, json!({"topic": "Ax runtime hooks"}), Value::Null, override_hooks)?);
        Ok(())
    })();
    set_rate_limiter(None);
    set_tracer(None);
    set_meter(None);
    result
}