16use std::process::ExitCode;

18use jev_protocol::{Choice, Json, Questions};
19
20
21const USAGE: &str = "usage: op-env-run -- jevprobe [--headers]
22  Asks Jev one Choice question about a made-up A Link to the Past position and
23  prints the whole distribution. Reads TYPESAFE_API_KEY from the environment.
24  --headers  print every response header instead of the distribution.";
25
26fn main() -> ExitCode {
27    // Piping this into `head` must not panic the process.
28    unsafe { libc::signal(libc::SIGPIPE, libc::SIG_DFL) };
29    let mut headers = false;
30    for arg in std::env::args().skip(1) {
31        match arg.as_str() {
32            "--headers" => headers = true,
33            _ => {
34                eprintln!("{USAGE}");
35                return ExitCode::from(2);
36            }
37        }
38    }
39    let runtime = match tokio::runtime::Builder::new_current_thread().enable_all().build() {
40        Ok(runtime) => runtime,
41        Err(e) => {
42            eprintln!("building the runtime: {e}");
43            return ExitCode::FAILURE;
44        }
45    };
46    match runtime.block_on(probe(headers)) {
47        Ok(()) => ExitCode::SUCCESS,
48        Err(e) => {
49            eprintln!("{e}");
50            ExitCode::FAILURE
51        }
52    }
53}

56fn json_of(value: &serde_json::Value) -> Json {
57    Json::verbatim(&value.to_string()).expect("serde_json writes JSON")
58}

60async fn probe(headers_only: bool) -> Result<(), String> {
61    let client = match jev_http::Jev::from_env("jevprobe", decisions::model()) {
62        Some(built) => built?,
63        None => {
64            return Err(format!(
65                "{} is not set. It comes from 1Password: `op-env-run -- jevprobe`.",
66                jev_http::KEY
67            ));
68        }
69    };
70
71    // A state shaped the way the real loop shapes one: an object with named
72    // parts, only what the question needs in it (digest §5).
73    let state = serde_json::json!({
74        "goal": "Leave Link's house and get outside.",
75        "where": {
76            "place": "indoors, Link's house, just out of bed",
77            "link": { "x": 2096, "y": 2000, "facing": "down" },
78        },
79        "around_link": {
80            "north": "bed and wall",
81            "south": "floor",
82            "east": "wall",
83            "west": "floor",
84        },
85        "landmarks": [
86            { "what": "staircase down and out", "direction": "south", "distance_px": 40 },
87        ],
88    });
89
90    let instructions = serde_json::json!({
91        "question": "What should Link do next?",
92        "focus": "Pick the ONE action that best advances `goal`. Code walks Link and \
93                  checks walls; you choose the direction.",
94    });
95    let question = Choice::new(
96        json_of(&instructions),
97        [
98            ("walk_north", "Walk up, away from the door"),
99            ("walk_south", "Walk down, toward the staircase"),
100            ("walk_west", "Walk left across the room"),
101            ("press_a", "Use what is directly in front of Link"),
102        ]
103        .map(|(id, what)| (id.to_owned(), Some(Json::text(what)))),
104    )
105    .map_err(|e| format!("the probe's question is not a valid Choice: {e}"))?;
106
107    let mut questions = Questions::new();
108    let mv = questions.choice("move", question).map_err(|e| e.to_string())?;
109    println!("asking {} …", client.model());
110    let answered = client
111        .ask(&json_of(&state), &questions, "jevprobe: one probe question")
112        .await
113        .map_err(|e| e.to_string())?;
114    let usage = answered.response.usage();
115
116    if headers_only {
117        println!();
118        println!("{} response header(s), anything credential-shaped redacted:", answered.headers.len());
119        let widest = answered.headers.iter().map(|(name, _)| name.len()).max().unwrap_or(0);
120        for (name, value) in &answered.headers {
121            println!("  {name:widest$}  {value}");
122        }
123        println!();
124        println!(
125            "request id  {}",
126            answered.request_id.as_deref().unwrap_or("(none)")
127        );
128        println!("rate limit  {:?}", answered.rate_limit);
129        println!("usage       {} in, {} out", usage.input_tokens, usage.output_tokens);
130        return Ok(());
131    }
132
133    let chosen = answered.response.get(mv);
134    println!();
135    println!("model      {}", answered.response.model());
136    println!("round trip {:.0} ms over {} attempt(s)", answered.took.as_secs_f64() * 1000.0, answered.attempts);
137    println!("tokens     {} in, {} out", usage.input_tokens, usage.output_tokens);
138    println!("cost       ${:.6}", usage.dollars());
139    println!("confidence {:.3}", chosen.confidence);
140    println!("chose      {}", chosen.choice);
141    println!();
142    let mut sorted: Vec<(&String, &f64)> = chosen.probabilities.iter().map(|(o, p)| (o, p)).collect();
143    sorted.sort_by(|a, b| b.1.total_cmp(a.1));
144    for (option, probability) in sorted {
145        let bar = "#".repeat((probability * 40.0).round() as usize);
146        println!("  {probability:6.3}  {bar:<40}  {option}");
147    }
148    println!();
149    println!("raw response:\n{}", answered.body);
150    Ok(())
151}