import dev.axllm.ax.*;
import java.nio.file.*;
import java.util.*;

public final class AxgenOptimizationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static AxAIService client() {
    return Ax.ai("openai",
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini"), "model_config", Map.of("temperature", 0.0)));
  }

  static final class ExampleOptimizer implements OptimizerEngine {
    public String name() { return "example"; }
    public String version() { return "1"; }
    public Map<String, Object> optimize(Map<String, Object> request) {
      return Map.of("componentMap", Map.of("priority::instruction", "Classify operational risk. Use high for production-impacting urgency."), "metadata", Map.of("source", "axgen"));
    }
  }

  public static void main(String[] args) throws Exception {
    AxGen program = new AxGen(Ax.s("emailText:string -> priority:class \"high, normal, low\", rationale:string"), Map.of("id", "priority", "instruction", "Classify the email priority."));
    Map<String, Object> baseline = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    Map<String, Object> artifact = program.optimizeWith(new ExampleOptimizer(), List.of(Map.of("emailText", "URGENT: checkout is down", "priority", "high")), Map.of("apply", false));
    program.applyOptimization(Json.stringify(artifact));
    Map<String, Object> after = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    System.out.println(Json.stringify(Map.of("baseline", baseline, "after", after)));
  }
}

import dev.axllm.ax.*;
import java.nio.file.*;
import java.util.*;

public final class GepaOptimizationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static AxAIService client() {
    return Ax.ai("openai",
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini"), "model_config", Map.of("temperature", 0.0)));
  }

  static final class LocalEvaluator implements OptimizerEvaluator {
    public Map<String, Object> evaluate(Map<String, Object> candidateMap, Map<String, Object> options) {
      return Map.of("rows", List.of(Map.of("prediction", Map.of("answer", "Ax composes typed LLM programs."), "scores", Map.of("quality", 0.9), "scalar", 0.9)), "avg", 0.9, "count", 1);
    }
  }

  public static void main(String[] args) throws Exception {
    AxGen program = new AxGen(Ax.s("emailText:string -> priority:class \"high, normal, low\", rationale:string"), Map.of("id", "priority", "instruction", "Classify the email priority."));
    Map<String, Object> baseline = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    Map<String, Object> request = Map.of("programKind", "axgen", "components", List.of(Map.of("id", "priority::instruction", "owner", "priority", "kind", "instruction", "current", "Classify priority clearly.")), "dataset", Map.of("train", List.of(Map.of("emailText", "URGENT: checkout is down"))), "options", Map.of("numTrials", 0, "maxMetricCalls", 4, "seed", 7));
    Map<String, Object> artifact = new AxGEPA(null, Map.of("seed", 7)).optimize(request, new LocalEvaluator());
    System.out.println(Json.stringify(Map.of("baseline", baseline, "artifact", artifact)));
  }
}

import dev.axllm.ax.*;
import java.nio.file.*;
import java.util.*;

public final class ArtifactOptimizationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static AxAIService client() {
    return Ax.ai("openai",
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini"), "model_config", Map.of("temperature", 0.0)));
  }

  static final class ExampleOptimizer implements OptimizerEngine {
    public String name() { return "example"; }
    public String version() { return "1"; }
    public Map<String, Object> optimize(Map<String, Object> request) {
      return Map.of("componentMap", Map.of("priority::instruction", "Classify operational risk. Use high for production-impacting urgency."), "metadata", Map.of("source", "artifact"));
    }
  }

  public static void main(String[] args) throws Exception {
    AxGen program = new AxGen(Ax.s("emailText:string -> priority:class \"high, normal, low\", rationale:string"), Map.of("id", "priority", "instruction", "Classify the email priority."));
    Map<String, Object> baseline = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    Map<String, Object> artifact = program.optimizeWith(new ExampleOptimizer(), List.of(Map.of("emailText", "URGENT: checkout is down", "priority", "high")), Map.of("apply", false));
    program.applyOptimization(Json.stringify(artifact));
    Map<String, Object> after = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    System.out.println(Json.stringify(Map.of("baseline", baseline, "after", after)));
  }
}

import dev.axllm.ax.*;
import dev.axllm.ax.runtime.quickjs.*;
import java.util.*;
import java.util.function.Consumer;

public final class AgentPlaybookEvolveExample {
  static String apiKey() {
    String value = System.getenv("OPENAI_API_KEY");
    if (value == null || value.isBlank()) value = System.getenv("OPENAI_APIKEY");
    if (value == null || value.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return value;
  }

  public static void main(String[] args) throws Exception {
    AxAIService client = Ax.ai("openai", Map.of(
        "api_key", apiKey(),
        "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-6-luna")));

    Map<String, Object> bullet = Map.of(
        "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");
    Map<String, Object> seed = Map.of(
        "playbook", Map.of(
            "version", 1,
            "sections", Map.of("failures_to_avoid", List.of(bullet)),
            "updatedAt", "2026-07-15T00:00:00.000Z"),
        "artifact", Map.of("feedback", List.of(), "history", List.of()));

    List<Object> observedCitations = new ArrayList<>();
    List<Object> playbookUpdates = new ArrayList<>();
    Consumer<List<Object>> citationObserver = observedCitations::add;
    Consumer<Map<String, Object>> playbookObserver = playbookUpdates::add;
    AxAgent assistant = Ax.agent(
        "question:string -> answer:string",
        Map.of(
            "ai", client,
            "contextFields", List.of(),
            "runtime", Map.of("language", "JavaScript"),
            "playbook", Map.of("playbook", seed, "seed", 7, "onUpdate", playbookObserver),
            "citations", Map.of("surface", "hidden", "onCitations", citationObserver)));
    assistant
        .setInstruction("Answer from evidence and state uncertainty plainly.")
        .addActorInstruction("Before finishing, verify the answer against the collected evidence.");

    try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
      Map<String, Object> answer = assistant.forward(
          client,
          Map.of("question", "What should a support agent verify before answering?"),
          Map.of("runtime", runtime, "max_actor_steps", 8));

      Map<String, Object> dataset = Map.of(
          "train", List.of(Map.of(
              "input", Map.of("question", "Give a concise evidence-first answer."),
              "score", 0)));
      Map<String, Object> evolution = assistant.playbook(null).evolve(
          dataset,
          Map.of("verify", true, "maxProposals", 1, "runtime", runtime));

      // Each train task replays on the runtime. A replay that fails scores 0
      // and would be mined as the weakness, so stop instead.
      List<Object> replays = new java.util.ArrayList<>();
      for (Object raw : (List<?>) evolution.get("records")) {
        Map<?, ?> record = (Map<?, ?>) raw;
        Map<?, ?> prediction = record.get("prediction") instanceof Map<?, ?> found ? found : Map.of();
        if (record.get("error") != null || "error".equals(prediction.get("completionType"))) {
          throw new IllegalStateException("evolve replay failed: " + (record.get("error") != null ? record.get("error") : prediction.get("error")));
        }
        replays.add(prediction.get("completionType"));
      }

      System.out.println(Json.pretty(answer));
      System.out.println("replays: " + replays);
      System.out.println("citations: " + (observedCitations.isEmpty() ? List.of() : observedCitations.get(observedCitations.size() - 1)));
      System.out.println("run-end updates: " + playbookUpdates.size());
      System.out.println("outcomes: " + evolution.get("outcomes"));
      System.out.println(assistant.getPlaybook().render());
    }
  }
}