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

struct ExampleOptimizer;
impl OptimizerEngine for ExampleOptimizer {
    fn optimize(&mut self, _request: Value, _evaluator: &mut dyn FnMut(Value) -> AxResult<Value>) -> AxResult<Value> {
        Ok(json!({"componentMap": {"priority::instruction": "Classify operational risk. Use high for production-impacting urgency."}, "metadata": {"source": "axgen"}}))
    }
}

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-6-luna".to_string());
    axllm::ai("openai", json!({"api_key": api_key, "model": model, "model_config": {"temperature": 0}}))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;
    let mut program = ax("emailText:string -> priority:class \"high, normal, low\", rationale:string")?;
    let baseline = program.forward(&mut client, json!({"emailText": "Production checkout is failing for enterprise customers."}))?;
    let mut optimizer = ExampleOptimizer;
    let artifact = optimizer.optimize(json!({"candidate": "priority"}), &mut |_candidate| Ok(json!({"score": 1.0})))?;
    let score = novel_f1_score("", baseline["rationale"].as_str().unwrap_or_default(), "Production checkout is failing for enterprise customers.", false)?;
    println!("{}", serde_json::to_string_pretty(&json!({"baseline": baseline, "artifact": artifact, "rationaleNovelF1": score}))?);
    Ok(())
}

use axllm::{ax, AxResult, OpenAICompatibleClient, OptimizerEngine};
use serde_json::{json, Value};
use std::{cell::RefCell, env, rc::Rc};

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-6-luna".to_string());
    axllm::ai("openai", json!({"api_key": api_key, "model": model}))
}
fn task(text: &str, priority: &str) -> Value {
    json!({"input": {"emailText": text}, "expectedOutput": {"priority": priority}})
}
fn main() -> AxResult<()> {
    let client = Rc::new(RefCell::new(openai_client()?));
    let mut program = ax("emailText:string -> priority:class \"high, normal, low\", rationale:string")?;
    let baseline = program.forward(&mut *client.borrow_mut(), json!({"emailText": "Production checkout is failing for enterprise customers."}))?;
    let request = json!({
        "programKind": "axgen", "components": program.get_optimizable_components(),
        "dataset": {"train": [task("URGENT: checkout is down", "high"), task("Weekly product newsletter", "low")],
            "validation": [task("Production payments are failing", "high"), task("Optional community meetup invite", "low")]},
        "options": {"numTrials": 1, "minibatchSize": 2, "maxMetricCalls": 12, "bootstrap": false, "seed": 7},
    });
    let mut engine = axllm::AxGEPA::new().with_reflection_client(client.clone());
    let artifact = engine.optimize(request, &mut |step| {
        let mut result = program.evaluate_optimization(&mut *client.borrow_mut(), &step["dataset"], &step["candidateMap"], &step["options"])?;
        let rows = result["rows"].as_array_mut().unwrap();
        let mut sum = 0.0;
        for row in rows.iter_mut() {
            let score = if row["prediction"]["output"]["priority"] == row["input"]["expectedOutput"]["priority"] { 1.0 } else { 0.0 };
            row["scores"] = json!({"accuracy": score});
            row["scalar"] = json!(score);
            sum += score;
        }
        let avg = sum / rows.len() as f64;
        result["avg"] = json!(avg);
        Ok(result)
    })?;
    let original = artifact["baselineScore"].as_f64().unwrap();
    let best = artifact["metadata"]["bestScore"].as_f64().unwrap();
    println!("{}", serde_json::to_string_pretty(&json!({"baselinePrediction": baseline, "baselineScore": original,
        "baselineScores": artifact["baselineScores"], "optimizedScore": best, "delta": best-original}))?);
    Ok(())
}

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

struct ExampleOptimizer;
impl OptimizerEngine for ExampleOptimizer {
    fn optimize(&mut self, _request: Value, _evaluator: &mut dyn FnMut(Value) -> AxResult<Value>) -> AxResult<Value> {
        Ok(json!({"componentMap": {"priority::instruction": "Classify operational risk. Use high for production-impacting urgency."}, "metadata": {"source": "artifact"}}))
    }
}

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());
    axllm::ai("openai", json!({"api_key": api_key, "model": model, "model_config": {"temperature": 0}}))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;
    let mut program = ax("emailText:string -> priority:class \"high, normal, low\", rationale:string")?;
    let baseline = program.forward(&mut client, json!({"emailText": "Production checkout is failing for enterprise customers."}))?;
    let mut optimizer = ExampleOptimizer;
    let artifact = optimizer.optimize(json!({"candidate": "priority"}), &mut |_candidate| Ok(json!({"score": 1.0})))?;
    println!("{}", serde_json::to_string_pretty(&json!({"baseline": baseline, "artifact": artifact}))?);
    Ok(())
}

use axllm::runtime::quickjs::QuickJsCodeRuntime;
use axllm::{agent_with_options, AxResult, OpenAICompatibleClient};
use serde_json::{json, Value};
use std::cell::RefCell;
use std::env;
use std::rc::Rc;

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-6-luna".to_string());
    axllm::ai("openai", json!({"api_key": api_key, "model": model}))
}

fn main() -> AxResult<()> {
    let seed = json!({
        "playbook": {
            "version": 1,
            "sections": {
                "failures_to_avoid": [{
                    "id": "failures-to-avoid-00001",
                    "section": "failures_to_avoid",
                    "content": "Check the available evidence before answering.",
                    "helpfulCount": 0,
                    "harmfulCount": 0,
                    "createdAt": "2026-07-15T00:00:00.000Z",
                    "updatedAt": "2026-07-15T00:00:00.000Z"
                }]
            },
            "updatedAt": "2026-07-15T00:00:00.000Z"
        },
        "artifact": {"feedback": [], "history": []}
    });

    let observed_citations = Rc::new(RefCell::new(Value::Array(Vec::new())));
    let playbook_updates = Rc::new(RefCell::new(Vec::<Value>::new()));
    let mut assistant = agent_with_options(
        "question:string -> answer:string",
        json!({
            "contextFields": [],
            "runtime": {"language": "JavaScript"},
            "playbook": {"playbook": seed, "seed": 7},
            "citations": {"surface": "hidden"}
        }),
    )?
    .with_runtime(Box::new(QuickJsCodeRuntime::new()))?;
    assistant
        .set_instruction("Answer from evidence and state uncertainty plainly.")?
        .add_actor_instruction(
            "Before finishing, verify the answer against the collected evidence.",
        )?;
    let citation_sink = observed_citations.clone();
    assistant.set_citations_observer(move |value| *citation_sink.borrow_mut() = value);
    let playbook_sink = playbook_updates.clone();
    assistant.set_playbook_observer(move |value| playbook_sink.borrow_mut().push(value));

    let student = Rc::new(RefCell::new(openai_client()?));
    let answer = {
        let mut client = student.borrow_mut();
        assistant.forward_with_options(
            &mut *client,
            json!({"question": "What should a support agent verify before answering?"}),
            json!({"max_actor_steps": 8}),
        )?
    };

    let mut playbook = assistant.playbook(
        student.clone(),
        None::<Rc<RefCell<OpenAICompatibleClient>>>,
        json!({}),
    )?;
    let dataset = json!({
        "train": [{
            "input": {"question": "Give a concise evidence-first answer."},
            "score": 0
        }]
    });
    // The playbook borrows `student` to reflect while the agent runs, so the
    // agent runs on its own client for the same model.
    let evolution = {
        let mut runner = openai_client()?;
        playbook.evolve_agent(
            &mut assistant,
            &mut runner,
            &dataset,
            &json!({"verify": true, "maxProposals": 1}),
        )?
    };

    // Each train task replays on the runtime. A replay that fails scores 0 and
    // would be mined as the weakness, so stop instead.
    let mut replays = Vec::new();
    for record in evolution["records"].as_array().cloned().unwrap_or_default() {
        let completion = record["prediction"]["completionType"].clone();
        if !record["error"].is_null() || completion == "error" {
            return Err(axllm::AxError::runtime(format!(
                "evolve replay failed: {} {}",
                record["error"], record["prediction"]["error"]
            )));
        }
        replays.push(completion);
    }

    println!("{}", serde_json::to_string_pretty(&answer)?);
    println!("replays: {}", serde_json::Value::Array(replays));
    println!("citations: {}", observed_citations.borrow());
    println!("run-end updates: {}", playbook_updates.borrow().len());
    println!("outcomes: {}", evolution["outcomes"]);
    println!("{}", playbook.render());
    Ok(())
}