use axllm::runtime::quickjs::QuickJsCodeRuntime;
use axllm::{agent_with_options, AxResult, GoogleGeminiClient};
use serde_json::{json, Value};
use std::env;

fn gemini_client() -> AxResult<GoogleGeminiClient> {
    let api_key = env::var("GOOGLE_APIKEY")
        .map_err(|_| axllm::AxError::runtime("Set GOOGLE_APIKEY to run this example."))?;
    let model = env::var("AX_GEMINI_MODEL").unwrap_or_else(|_| "gemini-3.5-flash".to_string());
    Ok(GoogleGeminiClient::new(api_key, model).with_profile("google-gemini"))
}

// ---------------------------------------------------------------------------
// Synthetic CloudWatch-style export -- generated large on purpose. Dumping these
// raw events into a prompt would blow the context window. The agent keeps them
// in its runtime (contextFields) and only the *evidence it extracts* ever
// reaches the model. Deterministic so the example is reproducible.
// ---------------------------------------------------------------------------
fn build_log_dump() -> Vec<Value> {
    // 2026-03-02 13:00:00Z, incremented by 2 seconds per index.
    const START_EPOCH: i64 = 1772456400; // 2026-03-02T13:00:00Z
    let mut events: Vec<Value> = Vec::new();

    let mut push = |i: i64, mut event: Value| {
        let secs = START_EPOCH + i * 2;
        // Format the absolute time as an ISO-8601 UTC stamp (no chrono dependency).
        event["timestamp"] = Value::String(iso_utc(secs));
        event["requestId"] = Value::String(format!("req-{}", 100000 + i));
        events.push(event);
    };

    for i in 0..1600i64 {
        // Routine, healthy traffic across the fleet.
        push(
            i,
            json!({"level": "INFO", "service": "gateway", "statusCode": 200, "latencyMs": 40 + (i % 30), "message": "route ok GET /checkout"}),
        );
        push(
            i,
            json!({"level": "INFO", "service": "search-api", "statusCode": 200, "latencyMs": 70 + (i % 50), "message": "query ok q=shoes"}),
        );

        // Window A: payments-gw upstream timeouts spill into checkout-api 502s for
        // enterprise tenants, with retry storms + pool exhaustion.
        if (300..520).contains(&i) {
            push(
                i,
                json!({"level": "ERROR", "service": "payments-gw", "statusCode": 504, "latencyMs": 10000, "tenantTier": "enterprise", "message": "upstream timeout calling acquirer (10s)"}),
            );
            push(
                i,
                json!({"level": "ERROR", "service": "checkout-api", "statusCode": 502, "tenantTier": "enterprise", "message": "bad gateway from svc-payments-gw"}),
            );
            if i % 3 == 0 {
                push(
                    i,
                    json!({"level": "WARN", "service": "payments-gw", "message": "connection pool exhausted (max=64) waiting=200+"}),
                );
                push(
                    i,
                    json!({"level": "WARN", "service": "checkout-api", "tenantTier": "enterprise", "message": "user-visible: \"Payment could not be processed\""}),
                );
            }
        }

        // Window B: the nightly catalog-cron pins CPU and search-api returns 429s.
        if (1000..1120).contains(&i) {
            push(
                i,
                json!({"level": "WARN", "service": "catalog-cron", "latencyMs": 0, "message": "rebuild step pinning CPU at 95% on shared node"}),
            );
            push(
                i,
                json!({"level": "ERROR", "service": "search-api", "statusCode": 429, "message": "rate limited: downstream catalog unavailable"}),
            );
        }
    }

    events
}

// Minimal UTC epoch-seconds -> "YYYY-MM-DDTHH:MM:SSZ" formatter (proleptic Gregorian).
fn iso_utc(epoch_secs: i64) -> String {
    let days = epoch_secs.div_euclid(86_400);
    let secs_of_day = epoch_secs.rem_euclid(86_400);
    let (hh, mm, ss) = (secs_of_day / 3600, (secs_of_day % 3600) / 60, secs_of_day % 60);
    // Convert days-since-epoch (1970-01-01) to a calendar date.
    let mut z = days + 719_468;
    let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
    z -= era * 146_097;
    let doe = z;
    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146_096) / 365;
    let year = yoe + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let day = doy - (153 * mp + 2) / 5 + 1;
    let month = if mp < 10 { mp + 3 } else { mp - 9 };
    let year = if month <= 2 { year + 1 } else { year };
    format!(
        "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
        year, month, day, hh, mm, ss
    )
}

fn main() -> AxResult<()> {
    let mut client = gemini_client()?;

    let logs = build_log_dump();
    println!(
        "Generated {} log events (kept out of the prompt).",
        logs.len()
    );

    // `with_runtime` attaches the embedded JS engine so the agent loop can run.
    let mut log_rlm = agent_with_options(
        "task:string, logs:json \"Raw CloudWatch export; keep this out of the prompt\" -> architecture:string[] \"Services and how they call each other\", findings:json[] \"Each: issue, count, window, evidence, impact\", overallHealth:string, nextActions:string[]",
        json!({
            // The export stays in the runtime; only extracted evidence reaches the model.
            "contextFields": ["logs"],
            "contextPolicy": {"preset": "lean", "budget": "balanced"},
            "maxRuntimeChars": 12000,
            "runtime": {"language": "JavaScript"},
        }),
    )?
    .with_runtime(Box::new(QuickJsCodeRuntime::new()))?;

    let report = log_rlm.forward_with_options(
        &mut client,
        json!({
            "logs": logs,
            "task": "Infer the service architecture from the logs alone. Then find repeated errors, throttles, retries, and bad user states -- with the affected time window, an occurrence count, and concrete log evidence for each.",
        }),
        json!({"max_actor_steps": 40}),
    )?;

    println!("\n=== Report ===");
    println!("{}", serde_json::to_string_pretty(&report)?);
    println!("\n=== Usage ===");
    println!("{}", serde_json::to_string_pretty(&log_rlm.get_usage())?);
    Ok(())
}

use axllm::runtime::quickjs::QuickJsCodeRuntime;
use axllm::{agent_with_options, AxResult, GoogleGeminiClient};
use serde_json::json;
use std::env;

fn gemini_client() -> AxResult<GoogleGeminiClient> {
    let api_key = env::var("GOOGLE_APIKEY")
        .map_err(|_| axllm::AxError::runtime("Set GOOGLE_APIKEY to run this example."))?;
    let model = env::var("AX_GEMINI_MODEL").unwrap_or_else(|_| "gemini-3.5-flash".to_string());
    Ok(GoogleGeminiClient::new(api_key, model).with_profile("google-gemini"))
}

struct Module {
    path: String,
    imports: Vec<&'static str>,
    writes: &'static str,
}

// ---------------------------------------------------------------------------
// A large module-dependency index for a monorepo. Each block is a record the
// agent must *search* to answer -- the answers cannot be guessed, only computed
// by filtering the index. Generated large so it would not fit comfortably in a
// prompt; it lives in contextFields and is queried from the runtime.
// ---------------------------------------------------------------------------
fn build_module_index() -> Vec<Module> {
    let mut modules: Vec<Module> = vec![
        Module { path: "packages/api/middleware/auth.ts".into(), imports: vec!["packages/shared"], writes: "-" },
        Module { path: "packages/api/middleware/rateLimit.ts".into(), imports: vec!["packages/db"], writes: "-" },
        Module { path: "packages/api/routes/checkout.ts".into(), imports: vec!["packages/api/middleware/auth.ts", "packages/services/orders/createOrder.ts", "packages/services/payments/charge.ts"], writes: "-" },
        Module { path: "packages/api/routes/search.ts".into(), imports: vec!["packages/api/middleware/auth.ts", "packages/services/catalog/searchCatalog.ts"], writes: "-" },
        Module { path: "packages/services/orders/createOrder.ts".into(), imports: vec!["packages/db", "packages/clients/bus"], writes: "orders" },
        Module { path: "packages/services/orders/orderRepo.ts".into(), imports: vec!["packages/db"], writes: "orders" },
        Module { path: "packages/services/payments/charge.ts".into(), imports: vec!["packages/clients/acquirer", "packages/db"], writes: "payments" },
        Module { path: "packages/services/payments/refund.ts".into(), imports: vec!["packages/clients/acquirer", "packages/db"], writes: "refunds" },
        Module { path: "packages/services/catalog/searchCatalog.ts".into(), imports: vec!["packages/db"], writes: "-" },
        Module { path: "packages/clients/acquirer/index.ts".into(), imports: vec!["packages/shared"], writes: "-" },
        Module { path: "packages/clients/bus/index.ts".into(), imports: vec!["packages/shared"], writes: "-" },
    ];
    // Filler modules so the index is genuinely large; some also depend on the acquirer.
    for i in 0..110 {
        let dep = if i % 4 == 0 { "packages/clients/acquirer" } else { "packages/db" };
        let writes = if i % 6 == 0 { "audit" } else { "-" };
        modules.push(Module {
            path: format!("packages/services/feature{}/handler.ts", i),
            imports: vec![dep, "packages/shared"],
            writes,
        });
    }
    modules
}

fn main() -> AxResult<()> {
    let mut client = gemini_client()?;

    let modules = build_module_index();
    let codebase_index = modules
        .iter()
        .map(|m| {
            format!(
                "PATH: {}\nIMPORTS: {}\nWRITES: {}",
                m.path,
                m.imports.join(", "),
                m.writes
            )
        })
        .collect::<Vec<_>>()
        .join("\n\n");
    println!(
        "Module index: {} records (kept out of the prompt).",
        modules.len()
    );

    // `with_runtime` attaches the embedded JS engine so the agent loop can run.
    let mut analyst = agent_with_options(
        "context:string, question:string -> answer:string, paths:string[] \"Exact PATH values from the index that answer the question\"",
        json!({
            "contextFields": ["context"],
            "contextPolicy": {"preset": "adaptive", "budget": "balanced"},
            "contextOptions": {
                "description": "The context is a module index of \"PATH / IMPORTS / WRITES\" records. Answer by filtering those records in code -- never guess. Return exact PATH values verbatim.",
            },
            // The Peek context map: small, persistent orientation reused across queries.
            "contextMap": {"maxChars": 1800, "infiniteEvolve": false, "evolveSteps": 1},
            "runtime": {"language": "JavaScript"},
        }),
    )?
    .with_runtime(Box::new(QuickJsCodeRuntime::new()))?;

    let questions = [
        "Which modules import 'packages/clients/acquirer'? Give the exact PATH values.",
        "Which modules write to the 'orders' table?",
        "What are the direct IMPORTS of packages/api/routes/checkout.ts?",
    ];

    for question in questions {
        let output = analyst.forward_with_options(
            &mut client,
            json!({"context": codebase_index, "question": question}),
            json!({"max_actor_steps": 24}),
        )?;
        let answer = output
            .get("answer")
            .and_then(|v| v.as_str())
            .unwrap_or("");
        let paths = output
            .get("paths")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .map(|p| p.as_str().map(String::from).unwrap_or_else(|| p.to_string()))
                    .collect::<Vec<_>>()
                    .join(", ")
            })
            .unwrap_or_default();
        println!("\nQ: {}", question);
        println!("A: {}", answer);
        println!("Paths: {}", paths);
    }

    println!("\nThe context map evolved on the first query and was reused for the rest.");
    Ok(())
}

use axllm::runtime::quickjs::QuickJsCodeRuntime;
use axllm::{agent_with_options, AxResult, GoogleGeminiClient};
use serde_json::{json, Value};
use std::env;
use std::sync::Arc;

fn gemini_client() -> AxResult<GoogleGeminiClient> {
    let api_key = env::var("GOOGLE_APIKEY")
        .map_err(|_| axllm::AxError::runtime("Set GOOGLE_APIKEY to run this example."))?;
    let model = env::var("AX_GEMINI_MODEL").unwrap_or_else(|_| "gemini-3.5-flash".to_string());
    Ok(GoogleGeminiClient::new(api_key, model).with_profile("google-gemini"))
}

const MONTHS: [&str; 12] = [
    "Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec",
];

#[derive(Clone)]
struct Row {
    region: String,
    product: String,
    month_index: usize,
    month: String,
    units: i64,
    revenue: i64,
    return_rate: f64,
}

// ---------------------------------------------------------------------------
// The "warehouse": a few hundred rows that live in the host process and are
// reachable only through tools. The model never sees the rows -- it queries
// them. Deterministic so the example is reproducible.
// ---------------------------------------------------------------------------
fn build_warehouse() -> Vec<Row> {
    let regions = [
        "North", "South", "East", "West", "Central", "NW", "NE", "SE",
    ];
    let products = ["Widget-A", "Widget-B", "Gadget-X", "Gadget-Y"];
    let mut rows: Vec<Row> = Vec::new();
    let mut seed: i64 = 7;
    let mut rand = || {
        seed = (seed.wrapping_mul(1103515245).wrapping_add(12345)) & 0x7FFF_FFFF;
        seed as f64 / 0x7FFF_FFFF as f64
    };

    for region in regions {
        for product in products {
            // A planted winner: East + Gadget-X grows far faster than the rest.
            let trend = if product == "Gadget-X" && region == "East" {
                90.0
            } else {
                25.0
            };
            for (m, month) in MONTHS.iter().enumerate() {
                let units = (400.0 + rand() * 1200.0 + m as f64 * trend).round() as i64;
                let price = if product.starts_with("Gadget") { 60 } else { 38 };
                let return_rate = round3(
                    0.01 + rand() * 0.05 + if product == "Widget-B" { 0.03 } else { 0.0 },
                );
                rows.push(Row {
                    region: region.to_string(),
                    product: product.to_string(),
                    month_index: m,
                    month: month.to_string(),
                    units,
                    revenue: units * price,
                    return_rate,
                });
            }
        }
    }
    rows
}

fn round3(value: f64) -> f64 {
    (value * 1000.0).round() / 1000.0
}

fn round4(value: f64) -> f64 {
    (value * 10000.0).round() / 10000.0
}

fn opt_str(p: &Value, key: &str) -> Option<String> {
    p.get(key)
        .and_then(|v| v.as_str())
        .filter(|s| !s.is_empty())
        .map(String::from)
}

fn main() -> AxResult<()> {
    let mut client = gemini_client()?;

    let warehouse = Arc::new(build_warehouse());
    println!(
        "Warehouse rows: {} (kept out of the prompt).",
        warehouse.len()
    );

    // The schema/data dictionary is large-ish and goes into contextFields so the
    // agent orients on column meaning + business rules without the doc entering the prompt.
    let schema = r#"TABLE sales (one row per region x product x month)

COLUMNS
  region       text   one of: North, South, East, West, Central, NW, NE, SE
  product      text   one of: Widget-A, Widget-B, Gadget-X, Gadget-Y
  month        text   Jan..Dec (calendar order; monthIndex 0..11)
  units        int    units sold that month
  revenue      int    integer dollars (units * unit price; Gadgets cost more)
  returnRate   float  fraction of units returned, 0..1

BUSINESS RULES
  - "Growth" = change in monthly revenue from Jan to Dec for a region+product.
  - A return rate above 0.05 (5%) is flagged for quality review.
  - Compare like-for-like: always group by region AND product, not either alone.

TOOLS AVAILABLE (call them, never invent figures)
  query  filter + aggregate a slice -> {matched, totalUnits, totalRevenue, avgReturnRate}
  top    rank a metric ("revenue"|"units") grouped by "product"|"region" -> [{key, value}]
  trend  monthly revenue series (Jan..Dec) for one region + product"#;

    // --- Host tool handlers over the warehouse (the model never sees the rows) ---
    let mut runtime = QuickJsCodeRuntime::new();

    let wh = warehouse.clone();
    runtime.register_callable("query", move |p: Value| {
        let region = opt_str(&p, "region");
        let product = opt_str(&p, "product");
        let month = opt_str(&p, "month");
        let rows: Vec<&Row> = wh
            .iter()
            .filter(|r| {
                region.as_deref().is_none_or(|v| r.region == v)
                    && product.as_deref().is_none_or(|v| r.product == v)
                    && month.as_deref().is_none_or(|v| r.month == v)
            })
            .collect();
        let total_units: i64 = rows.iter().map(|r| r.units).sum();
        let total_revenue: i64 = rows.iter().map(|r| r.revenue).sum();
        let avg_return = if rows.is_empty() {
            0.0
        } else {
            round4(rows.iter().map(|r| r.return_rate).sum::<f64>() / rows.len() as f64)
        };
        Ok(json!({
            "matched": rows.len(),
            "totalUnits": total_units,
            "totalRevenue": total_revenue,
            "avgReturnRate": avg_return,
        }))
    })?;

    let wh = warehouse.clone();
    runtime.register_callable("top", move |p: Value| {
        let metric = opt_str(&p, "metric").unwrap_or_else(|| "revenue".to_string());
        let group_by = opt_str(&p, "groupBy").unwrap_or_else(|| "product".to_string());
        let limit = p.get("limit").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
        let mut totals: std::collections::HashMap<String, i64> = std::collections::HashMap::new();
        for r in wh.iter() {
            let key = if group_by == "region" {
                r.region.clone()
            } else {
                r.product.clone()
            };
            let value = if metric == "units" { r.units } else { r.revenue };
            *totals.entry(key).or_insert(0) += value;
        }
        let mut ranked: Vec<(String, i64)> = totals.into_iter().collect();
        ranked.sort_by(|a, b| b.1.cmp(&a.1));
        let out: Vec<Value> = ranked
            .into_iter()
            .take(limit)
            .map(|(key, value)| json!({"key": key, "value": value}))
            .collect();
        Ok(Value::Array(out))
    })?;

    let wh = warehouse.clone();
    runtime.register_callable("trend", move |p: Value| {
        let region = opt_str(&p, "region").unwrap_or_default();
        let product = opt_str(&p, "product").unwrap_or_default();
        let mut series = [0i64; 12];
        for r in wh.iter() {
            if r.region == region && r.product == product {
                series[r.month_index] = r.revenue;
            }
        }
        Ok(json!(series.to_vec()))
    })?;

    // `with_runtime` attaches the embedded JS engine (with the tools above) so
    // the agent loop can run and call them.
    let mut analyst = agent_with_options(
        "schema:string, question:string -> answer:string, evidence:string[] \"Concrete figures the answer is based on\"",
        json!({
            // Big data dictionary stays out of the prompt.
            "contextFields": ["schema"],
            // Tool specs advertised to the model; handlers are registered on the runtime above.
            "functions": [
                {
                    "name": "query",
                    "description": "Filter the sales table and return aggregates for the matching rows.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "region": {"type": "string"},
                            "product": {"type": "string"},
                            "month": {"type": "string"},
                        },
                    },
                },
                {
                    "name": "top",
                    "description": "Rank a metric (revenue|units) grouped by product|region, highest first.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "metric": {"type": "string"},
                            "groupBy": {"type": "string"},
                            "limit": {"type": "number"},
                        },
                        "required": ["metric", "groupBy"],
                    },
                },
                {
                    "name": "trend",
                    "description": "Monthly revenue series (Jan..Dec) for one region and product.",
                    "parameters": {
                        "type": "object",
                        "properties": {"region": {"type": "string"}, "product": {"type": "string"}},
                        "required": ["region", "product"],
                    },
                },
            ],
            "contextPolicy": {"preset": "lean", "budget": "balanced"},
            "runtime": {"language": "JavaScript"},
        }),
    )?
    .with_runtime(Box::new(runtime))?;

    let result = analyst.forward_with_options(
        &mut client,
        json!({
            "schema": schema,
            "question": "Which region+product had the strongest Jan->Dec revenue growth, and which products have an average return rate above the 5% review threshold?",
        }),
        json!({"max_actor_steps": 40}),
    )?;

    println!("{}", serde_json::to_string_pretty(&result)?);
    Ok(())
}

use axllm::runtime::quickjs::QuickJsCodeRuntime;
use axllm::{agent_with_options, ax, AxResult, OpenAICompatibleClient};
use serde_json::{json, Value};
use std::collections::BTreeMap;
use std::env;
use std::sync::{Arc, Mutex};

fn openai_config() -> AxResult<(String, String)> {
    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((api_key, model))
}

fn openai_client(api_key: &str, model: &str) -> OpenAICompatibleClient {
    OpenAICompatibleClient::new(api_key.to_string(), model.to_string())
        .with_model_config(json!({"temperature": 0}))
}

// ---------------------------------------------------------------------------
// The "lab": a deterministic black-box experiment. It scores an ETL config plan
// against a hidden ideal and returns, for any failing check, the exact fix --
// so the agent can converge by following the feedback, not by being told.
// ---------------------------------------------------------------------------
const CHECKS: [&str; 5] = [
    "no-nulls",
    "no-duplicates",
    "numeric-types",
    "trimmed-strings",
    "outliers-handled",
];

fn remedy(check: &str) -> &'static str {
    match check {
        "no-nulls" => "set nullPolicy=impute (or nullPolicy=drop)",
        "no-duplicates" => "set dedup=on",
        "numeric-types" => "set coerceTypes=on",
        "trimmed-strings" => "set trim=on",
        "outliers-handled" => "set outlier=clip (or outlier=winsorize)",
        _ => "",
    }
}

// Parse `key=value` flags (lowercased) out of a free-form plan string, matching
// the Python reference's regex `([a-z]+)\s*=\s*([a-z0-9]+)`.
fn parse_flags(plan: &str) -> BTreeMap<String, String> {
    let lower = plan.to_lowercase();
    let mut flags = BTreeMap::new();
    let bytes = lower.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i].is_ascii_lowercase() {
            let key_start = i;
            while i < bytes.len() && bytes[i].is_ascii_lowercase() {
                i += 1;
            }
            let key = &lower[key_start..i];
            let mut j = i;
            while j < bytes.len() && (bytes[j] == b' ' || bytes[j] == b'\t') {
                j += 1;
            }
            if j < bytes.len() && bytes[j] == b'=' {
                j += 1;
                while j < bytes.len() && (bytes[j] == b' ' || bytes[j] == b'\t') {
                    j += 1;
                }
                let val_start = j;
                while j < bytes.len() && bytes[j].is_ascii_alphanumeric() {
                    j += 1;
                }
                if j > val_start {
                    flags.insert(key.to_string(), lower[val_start..j].to_string());
                    i = j;
                    continue;
                }
            }
        } else {
            i += 1;
        }
    }
    flags
}

fn run_in_sandbox(plan: &str) -> Value {
    let flags = parse_flags(plan);
    let is = |k: &str, opts: &[&str]| flags.get(k).is_some_and(|v| opts.contains(&v.as_str()));
    let ok = |check: &str| match check {
        "no-nulls" => is("nullpolicy", &["impute", "drop"]),
        "no-duplicates" => is("dedup", &["on"]),
        "numeric-types" => is("coercetypes", &["on"]),
        "trimmed-strings" => is("trim", &["on"]),
        "outliers-handled" => is("outlier", &["clip", "winsorize"]),
        _ => false,
    };
    let passed: Vec<&str> = CHECKS.into_iter().filter(|c| ok(c)).collect();
    let failed: Vec<Value> = CHECKS
        .into_iter()
        .filter(|c| !ok(c))
        .map(|c| json!({"check": c, "fix": remedy(c)}))
        .collect();
    let score = ((passed.len() as f64 / CHECKS.len() as f64) * 100.0).round() / 100.0;
    json!({
        "score": score,
        "solved": passed.len() == CHECKS.len(),
        "passed": passed,
        "failed": failed,
        "logs": format!("{}/{} checks passed", passed.len(), CHECKS.len()),
    })
}

fn main() -> AxResult<()> {
    let (api_key, model) = openai_config()?;
    let mut client = openai_client(&api_key, &model);

    // In-memory rule store. Verified, reusable rules go here -- not raw failure notes.
    let memory_store: Arc<Mutex<BTreeMap<String, String>>> = Arc::new(Mutex::new(BTreeMap::new()));

    let mut runtime = QuickJsCodeRuntime::new();

    runtime.register_callable("runExperiment", |p: Value| {
        let plan = p.get("plan").and_then(|v| v.as_str()).unwrap_or("");
        Ok(run_in_sandbox(plan))
    })?;

    runtime.register_callable("listChecks", |_p: Value| Ok(json!(CHECKS.to_vec())))?;

    // An independent verifier -- a separate ax() program, not the agent grading
    // itself. It runs on its own client built from the same credentials so the
    // grade callable stays Send + Sync.
    let grade_key = api_key.clone();
    let grade_model = model.clone();
    runtime.register_callable("grade", move |p: Value| {
        let mut verifier = ax(
            "rubric:string, evidence:json -> passed:boolean, feedback:string, missing:string[]",
        )?;
        verifier.set_instruction(
            "You are an independent rubric grader, not a self-critique. Pass only when the evidence clearly satisfies every part of the rubric.",
        );
        let mut grader = openai_client(&grade_key, &grade_model);
        verifier.forward(
            &mut grader,
            json!({
                "rubric": p.get("rubric").cloned().unwrap_or_else(|| json!("")),
                "evidence": p.get("evidence").cloned().unwrap_or_else(|| json!([])),
            }),
        )
    })?;

    // NOTE: `recall` is a reserved runtime builtin in the QuickJS engine, so the
    // host tool is named `recallRules` (the model is told to use that name).
    let recall_store = memory_store.clone();
    runtime.register_callable("recallRules", move |p: Value| {
        let topic = p
            .get("topic")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_lowercase();
        let words: Vec<&str> = topic.split_whitespace().collect();
        let store = recall_store.lock().unwrap();
        let hits: Vec<Value> = store
            .iter()
            .filter(|(k, _)| k.contains(&topic) || words.iter().any(|w| k.contains(w)))
            .map(|(_, v)| json!(v))
            .collect();
        Ok(Value::Array(hits))
    })?;

    let remember_store = memory_store.clone();
    runtime.register_callable("remember", move |p: Value| {
        let rule = p.get("rule").and_then(|v| v.as_str()).unwrap_or("").to_string();
        let evidence = p
            .get("evidence")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string();
        let key: String = rule.to_lowercase().chars().take(48).collect();
        let mut store = remember_store.lock().unwrap();
        store.insert(key, format!("{rule} :: {evidence}"));
        Ok(json!({"stored": true, "total": store.len()}))
    })?;

    let executor_description = [
        "Use the tools -- do not answer from your own knowledge.",
        "1. recallRules('etl data quality') to reuse anything already learned.",
        "2. runExperiment('') once to see every failing check and its fix.",
        "3. Build a plan applying all the fixes, then runExperiment again. Repeat until solved is true.",
        "4. grade the passing evidence against the rubric.",
        "5. For each check you fixed, remember(rule, evidence).",
        "6. Then return the answer, the plans you tried, and the learned rules.",
    ]
    .join("\n");

    let mut self_improving = agent_with_options(
        "goal:string, rubric:string -> answer:string, experiments:string[] \"Plans tried, in order\", learnedRules:string[]",
        json!({
            "contextFields": [],
            "functions": [
                {
                    "name": "runExperiment",
                    "description": "Apply an ETL config plan; returns score, solved, passed[], failed[{check,fix}], logs. Pass an empty plan to discover the fixes.",
                    "parameters": {"type": "object", "properties": {"plan": {"type": "string"}}, "required": ["plan"]},
                },
                {
                    "name": "listChecks",
                    "description": "List the data-quality checks the experiment evaluates.",
                    "parameters": {"type": "object", "properties": {}},
                },
                {
                    "name": "grade",
                    "description": "Independent rubric grader. Pass only when the evidence meets the rubric.",
                    "parameters": {"type": "object", "properties": {"rubric": {"type": "string"}, "evidence": {"type": "array", "items": {"type": "string"}}}, "required": ["rubric", "evidence"]},
                },
                {
                    "name": "recallRules",
                    "description": "Recall verified rules relevant to a topic.",
                    "parameters": {"type": "object", "properties": {"topic": {"type": "string"}}, "required": ["topic"]},
                },
                {
                    "name": "remember",
                    "description": "Store a verified, reusable rule (the rule, not raw notes).",
                    "parameters": {"type": "object", "properties": {"rule": {"type": "string"}, "evidence": {"type": "string"}}, "required": ["rule", "evidence"]},
                },
            ],
            "contextPolicy": {"preset": "adaptive", "budget": "balanced"},
            "executorOptions": {"description": executor_description},
            "runtime": {"language": "JavaScript"},
        }),
    )?
    .with_runtime(Box::new(runtime))?;

    let result = self_improving.forward_with_options(
        &mut client,
        json!({
            "goal": "Find an ETL config plan that cleans the dirty dataset so every data-quality check passes.",
            "rubric": "All five checks (no-nulls, no-duplicates, numeric-types, trimmed-strings, outliers-handled) must pass, i.e. score 1.0.",
        }),
        json!({"max_actor_steps": 18}),
    )?;

    println!("{}", serde_json::to_string_pretty(&result)?);

    // Persist the agent's verified rules so a future run's recall reuses them.
    if let Some(rules) = result.get("learnedRules").and_then(|v| v.as_array()) {
        let mut store = memory_store.lock().unwrap();
        for rule in rules {
            let rule_str = rule.as_str().map(String::from).unwrap_or_else(|| rule.to_string());
            let key: String = rule_str.to_lowercase().chars().take(48).collect();
            store.insert(key, rule_str);
        }
    }
    let total = memory_store.lock().unwrap().len();
    println!("\nMemory now holds {total} rule(s) for next time.");
    Ok(())
}