synthiaresearch 0.1.2

//! The Synthia client. Mirrors `packages/sdk-js/src/client.ts` (canonical) and
//! `packages/sdk-python/src/synthia/client.py` method-for-method.
//!
//! Deliberate Rust divergences, all listed in `docs/reference/rust-api.md`:
//! async on tokio (JS is async, Python is blocking + a thread pool), `Result`
//! instead of exceptions, and options structs instead of keyword arguments.
//!
//! Request bodies are `serde_json::json!` literals at the call site, never
//! derived `Serialize` structs: a struct would send `"seed_ids": []` where JS
//! and Python omit the key, and `scripts/check-request-coverage.py` reads the
//! keys lexically next to the path string.

use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use rand::Rng;
use serde_json::{json, Map, Value};
use tokio::sync::Mutex;

use crate::error::{Error, Result};
use crate::sandbox::{ToolEvent, ToolSandbox};
use crate::types;
use crate::{DEFAULT_BASE_URL, SDK_VERSION};

/// Transient server statuses worth retrying (cold DB, overload, gateway).
pub(crate) const RETRYABLE_STATUS: &[u16] = &[500, 502, 503, 504, 429, 529];

/// Create-style POSTs where a retried duplicate is harmless because the caller
/// keys off the returned id. Excludes turn-advance POSTs, which must never be
/// replayed — a re-sent reply could double-advance the server transcript.
const CREATE_POSTS: &[&str] = &["/v1/sdk-sessions", "/v1/rollouts", "/v1/quality-checks"];

/// The server bounds one quality check's LLM fan-out; bigger runs judge in
/// chunks and pool the evaluations.
const QUALITY_CHECK_CHUNK: usize = 500;

/// Extensions that make a reply or seed body a file upload. Kept identical to
/// the JS `FILE_SUFFIXES` and Python `_FILE_SUFFIXES`.
const FILE_SUFFIXES: &[&str] = &[
    "wav", "mp3", "m4a", "ogg", "flac", "webm", "aac", "aiff", "png", "jpg", "jpeg", "gif", "webp",
    "pdf", "docx", "pptx", "xlsx",
];

/// Backoff with jitter: 0.5s, 1s, 2s (+ up to 250ms). Same curve as the peers.
fn backoff(attempt: u32) -> Duration {
    let jitter = rand::thread_rng().gen_range(0..250);
    Duration::from_millis(500 * 2u64.pow(attempt) + jitter)
}

fn is_create_post(method: &str, path: &str) -> bool {
    method == "POST" && CREATE_POSTS.contains(&path)
}

// --- agent contracts --------------------------------------------------------

/// One traced tool invocation made by the agent while answering a probe.
#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
pub struct ToolCall {
    pub name: String,
    pub input: Map<String, Value>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub output: Option<Map<String, Value>>,
    #[serde(default)]
    pub is_error: bool,
}

/// An agent's reply plus the tool calls it made producing it.
#[derive(Clone, Debug, Default)]
pub struct AgentReply {
    pub reply: String,
    pub tool_calls: Vec<ToolCall>,
}

impl From<String> for AgentReply {
    fn from(reply: String) -> Self {
        Self {
            reply,
            tool_calls: Vec::new(),
        }
    }
}

impl From<&str> for AgentReply {
    fn from(reply: &str) -> Self {
        Self::from(reply.to_string())
    }
}

/// A rollout agent's reply: text, or a file (bytes / a path).
///
/// Audio replies flip the rollout into voice mode (the server talks back in
/// audio); images and documents are rendered to text. Either way the
/// server-extracted text drives the simulator.
#[derive(Clone, Debug)]
pub enum Reply {
    Text(String),
    Bytes(Vec<u8>),
    File(PathBuf),
}

impl From<String> for Reply {
    fn from(s: String) -> Self {
        Reply::Text(s)
    }
}
impl From<&str> for Reply {
    fn from(s: &str) -> Self {
        Reply::Text(s.to_string())
    }
}
impl From<Vec<u8>> for Reply {
    fn from(b: Vec<u8>) -> Self {
        Reply::Bytes(b)
    }
}
impl From<PathBuf> for Reply {
    fn from(p: PathBuf) -> Self {
        Reply::File(p)
    }
}

/// Anything that names a user model: an id, or the model itself.
///
/// The peers take `UserModel | string` on `generate`, `compose`, and
/// `submit_scenarios`; this is the Rust spelling of that union.
pub trait ModelRef {
    fn model_id(&self) -> &str;
}

impl ModelRef for str {
    fn model_id(&self) -> &str {
        self
    }
}
impl ModelRef for String {
    fn model_id(&self) -> &str {
        self
    }
}
impl ModelRef for types::UserModel {
    fn model_id(&self) -> &str {
        &self.id
    }
}
impl<T: ModelRef + ?Sized> ModelRef for &T {
    fn model_id(&self) -> &str {
        (**self).model_id()
    }
}

/// Anything that names a finished rollout: an id, or the result itself.
/// The peers take `(RolloutResult | string)[]` on `quality_check`.
pub trait RolloutRef {
    fn rollout_id(&self) -> &str;
}

impl RolloutRef for str {
    fn rollout_id(&self) -> &str {
        self
    }
}
impl RolloutRef for String {
    fn rollout_id(&self) -> &str {
        self
    }
}
impl RolloutRef for RolloutResult {
    fn rollout_id(&self) -> &str {
        &self.rollout_id
    }
}
impl<T: RolloutRef + ?Sized> RolloutRef for &T {
    fn rollout_id(&self) -> &str {
        (**self).rollout_id()
    }
}

/// A probe agent: answers one interview question at a time.
#[async_trait::async_trait]
pub trait ProbeAgent: Send + Sync {
    async fn answer(&self, probe: String) -> Result<AgentReply>;
}

/// A rollout agent: given the transcript so far and a deterministic
/// [`ToolSandbox`], produce the next reply.
///
/// Scenarios may run concurrently, so implementors must not share mutable
/// state across invocations — per-conversation state belongs on
/// `sandbox.context`, which is created fresh per rollout.
#[async_trait::async_trait]
pub trait RolloutAgent: Send + Sync {
    async fn reply(
        &self,
        transcript: Vec<types::TranscriptTurn>,
        sandbox: &mut ToolSandbox,
    ) -> Result<Reply>;
}

/// Drive a [`RolloutAgent`] with a probe question as a one-turn conversation,
/// tracing its sandbox calls onto the reply so probing still observes tool use.
struct ProbeFromRollout<'a>(&'a dyn RolloutAgent);

#[async_trait::async_trait]
impl ProbeAgent for ProbeFromRollout<'_> {
    async fn answer(&self, probe: String) -> Result<AgentReply> {
        let mut sandbox = ToolSandbox::new(0);
        let turn = types::TranscriptTurn {
            role: types::TranscriptTurnRole::User,
            content: probe,
            audio_url: None,
        };
        let reply = self.0.reply(vec![turn], &mut sandbox).await?;
        let text = match reply {
            Reply::Text(t) => t,
            _ => {
                return Err(Error::Invalid(
                    "probing needs a text reply — pass probe_agent for audio agents".to_string(),
                ))
            }
        };
        Ok(AgentReply {
            reply: text,
            tool_calls: sandbox
                .events
                .into_iter()
                .map(|e: ToolEvent| ToolCall {
                    name: e.name,
                    input: e.input,
                    output: Some(e.output),
                    is_error: e.is_error,
                })
                .collect(),
        })
    }
}

// --- file uploads -----------------------------------------------------------

struct Upload {
    file_b64: String,
    filename: Option<String>,
}

fn as_file_upload(reply: &Reply) -> Result<Option<Upload>> {
    use base64::Engine;
    match reply {
        Reply::Text(_) => Ok(None),
        Reply::Bytes(b) => Ok(Some(Upload {
            file_b64: base64::engine::general_purpose::STANDARD.encode(b),
            filename: None,
        })),
        Reply::File(p) => {
            let ext = p
                .extension()
                .and_then(|e| e.to_str())
                .map(str::to_lowercase)
                .unwrap_or_default();
            if !FILE_SUFFIXES.contains(&ext.as_str()) {
                return Ok(None);
            }
            let bytes = std::fs::read(p)?;
            Ok(Some(Upload {
                file_b64: base64::engine::general_purpose::STANDARD.encode(bytes),
                filename: p.file_name().and_then(|n| n.to_str()).map(str::to_string),
            }))
        }
    }
}

/// Seed content: JSON ingested as-is, or a file handled server-side with zero
/// parameters (audio transcribed and marked voice-origin, everything else
/// rendered to text).
#[derive(Clone, Debug)]
pub enum Content {
    Json(Value),
    Bytes(Vec<u8>),
    File(PathBuf),
}

impl Content {
    pub fn file<P: Into<PathBuf>>(p: P) -> Self {
        Content::File(p.into())
    }
}

impl From<Value> for Content {
    fn from(v: Value) -> Self {
        Content::Json(v)
    }
}

// --- HTTP -------------------------------------------------------------------

impl std::fmt::Debug for Http {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // Opaque on purpose: never render headers, which carry the API key.
        f.debug_struct("Http")
            .field("base_url", &self.base_url)
            .finish()
    }
}

struct Http {
    base_url: String,
    client: reqwest::Client,
    /// Auth plus the session headers added after the handshake.
    headers: Mutex<Vec<(String, String)>>,
}

impl Http {
    fn new(base_url: String, api_key: Option<String>) -> Result<Self> {
        let client = reqwest::Client::builder()
            // Generous: probe convergence runs model inference server-side.
            .timeout(Duration::from_secs(300))
            // Manual redirects: reqwest would auto-follow a 303 as a GET,
            // silently breaking POST semantics.
            .redirect(reqwest::redirect::Policy::none())
            .build()?;
        let mut headers = Vec::new();
        if let Some(key) = api_key {
            // Keyless is allowed (anonymous session), so only set it when present.
            headers.push(("authorization".to_string(), format!("Bearer {key}")));
        }
        Ok(Self {
            base_url,
            client,
            headers: Mutex::new(headers),
        })
    }

    async fn set_header(&self, name: &str, value: &str) {
        let mut headers = self.headers.lock().await;
        headers.retain(|(k, _)| k != name);
        headers.push((name.to_string(), value.to_string()));
    }

    async fn raw(
        &self,
        method: &str,
        path: &str,
        body: Option<&Value>,
        params: &[(&str, String)],
    ) -> Result<reqwest::Response> {
        let headers = self.headers.lock().await.clone();
        let url = format!("{}{}", self.base_url, path);

        let send = || async {
            let mut req = self
                .client
                .request(
                    reqwest::Method::from_bytes(method.as_bytes()).unwrap(),
                    &url,
                )
                .header("content-type", "application/json");
            for (k, v) in &headers {
                req = req.header(k, v);
            }
            if !params.is_empty() {
                req = req.query(params);
            }
            if let Some(b) = body {
                req = req.json(b);
            }
            let mut resp = req.send().await?;
            // Modal serves 303 attempt-token redirects while a deployment hands
            // requests between containers. The token resumes the SAME attempt,
            // and the async-poll pattern means a GET awaits the already-accepted
            // request's result — re-POSTing the body is rejected with a 400.
            let mut hop = 0;
            while resp.status().as_u16() == 303 && hop < 30 {
                let location = resp
                    .headers()
                    .get("location")
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or("")
                    .to_string();
                if !location.contains("__modal_attempt_token") {
                    break;
                }
                let mut req = self.client.get(&location);
                for (k, v) in &headers {
                    req = req.header(k, v);
                }
                resp = req.send().await?;
                hop += 1;
            }
            Ok::<reqwest::Response, Error>(resp)
        };

        // A minutes-long CI run is many round trips against serverless infra
        // that can transiently blip. One blip shouldn't fail the whole run.
        let retryable = method == "GET" || is_create_post(method, path);
        let max_attempts: u32 = if retryable { 4 } else { 1 };
        let mut last_err: Option<Error> = None;
        for attempt in 0..max_attempts {
            match send().await {
                Ok(resp) => {
                    if retryable
                        && RETRYABLE_STATUS.contains(&resp.status().as_u16())
                        && attempt < max_attempts - 1
                    {
                        tokio::time::sleep(backoff(attempt)).await;
                        continue;
                    }
                    return Ok(resp);
                }
                Err(e) => {
                    last_err = Some(e);
                    if attempt < max_attempts - 1 {
                        tokio::time::sleep(backoff(attempt)).await;
                    }
                }
            }
        }
        Err(last_err.unwrap_or_else(|| Error::Api("request failed".to_string())))
    }

    async fn check(resp: reqwest::Response, method: &str, path: &str) -> Result<reqwest::Response> {
        let status = resp.status().as_u16();
        if (200..300).contains(&status) {
            return Ok(resp);
        }
        Err(Error::Http {
            status,
            body: resp.text().await.unwrap_or_default(),
            method: method.to_string(),
            path: path.to_string(),
        })
    }

    async fn get<T: serde::de::DeserializeOwned>(
        &self,
        path: &str,
        params: &[(&str, String)],
    ) -> Result<T> {
        let resp = self.raw("GET", path, None, params).await?;
        let resp = Self::check(resp, "GET", path).await?;
        Ok(serde_json::from_str(&resp.text().await?)?)
    }

    async fn post<T: serde::de::DeserializeOwned>(&self, path: &str, body: Value) -> Result<T> {
        let resp = self.raw("POST", path, Some(&body), &[]).await?;
        let resp = Self::check(resp, "POST", path).await?;
        Ok(serde_json::from_str(&resp.text().await?)?)
    }

    /// Binary GET (audio artifacts).
    async fn get_bytes(&self, path: &str) -> Result<Vec<u8>> {
        let resp = self.raw("GET", path, None, &[]).await?;
        let resp = Self::check(resp, "GET", path).await?;
        Ok(resp.bytes().await?.to_vec())
    }
}

/// Cursor-based poller that prints a run's server-side telemetry.
struct EventStream {
    http: Arc<Http>,
    path: String,
    enabled: bool,
    after: i64,
}

impl EventStream {
    fn new(http: Arc<Http>, path: String, enabled: bool) -> Self {
        Self {
            http,
            path,
            enabled,
            after: 0,
        }
    }

    async fn pump(&mut self) -> Result<()> {
        if !self.enabled {
            return Ok(());
        }
        let body: types::EventList = self
            .http
            .get(&self.path, &[("after", self.after.to_string())])
            .await?;
        for event in body.data {
            self.after = event.seq;
            let detail = event
                .data
                .iter()
                .map(|(k, v)| format!("{k}={}", render_scalar(v)))
                .collect::<Vec<_>>()
                .join(" ");
            let suffix = if detail.is_empty() {
                String::new()
            } else {
                format!("  ({detail})")
            };
            println!("    ~ [{}] {}{}", event.stage, event.message, suffix);
        }
        Ok(())
    }
}

/// Render a JSON scalar the way JS string interpolation and Python f-strings do
/// for telemetry details — bare, not JSON-quoted.
fn render_scalar(v: &Value) -> String {
    match v {
        Value::String(s) => s.clone(),
        other => other.to_string(),
    }
}

/// The wire value of a transcript turn's role. The peers expose `role` as a
/// plain string; here the generated type is an open enum, so this is the
/// equivalent accessor (`"user"`, `"agent"`, or an unknown value verbatim).
pub fn role_str(role: &types::TranscriptTurnRole) -> String {
    match role {
        types::TranscriptTurnRole::User => "user".to_string(),
        types::TranscriptTurnRole::Agent => "agent".to_string(),
        types::TranscriptTurnRole::Other(s) => s.clone(),
    }
}

fn pct(x: f64) -> String {
    format!("{}%", (x * 100.0).round() as i64)
}

fn same_id_set(a: &[String], b: &[String]) -> bool {
    let sa: HashSet<&String> = a.iter().collect();
    let sb: HashSet<&String> = b.iter().collect();
    sa == sb
}

// --- options ----------------------------------------------------------------

#[derive(Clone, Debug, Default)]
pub struct SynthiaOptions {
    pub api_key: Option<String>,
    pub base_url: Option<String>,
    /// Session name. `Some(name)` pins it; `None` resolves from
    /// `SYNTHIA_SESSION` then a derived "project/script" name.
    pub session: Option<String>,
    /// Opt out into a fresh ephemeral session (the JS `session: false`).
    pub ephemeral_session: bool,
    /// CI provenance for this process, sent on the handshake and stamped onto
    /// every run this invocation creates. Set by `synthia run`; reporting only.
    pub ci: Option<Value>,
}

#[derive(Clone, Debug)]
pub struct WaitOptions {
    pub poll_interval: f64,
    pub timeout: f64,
    pub verbose: bool,
}

impl Default for WaitOptions {
    fn default() -> Self {
        Self {
            poll_interval: 2.0,
            timeout: 1800.0,
            verbose: false,
        }
    }
}

#[derive(Clone, Debug)]
pub struct ProbeOptions {
    pub max_turns: i64,
    pub verbose: bool,
    /// Already-ingested seeds to fold into the resulting model. `None` leaves
    /// grounding unspecified; `Some(vec![])` explicitly asks for ungrounded.
    pub seed_ids: Option<Vec<String>>,
}

impl Default for ProbeOptions {
    fn default() -> Self {
        Self {
            max_turns: 10,
            verbose: false,
            seed_ids: None,
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct GenerateOptions {
    /// Defaults to 20, matching the peers.
    pub count: Option<i64>,
    pub quality_check_id: Option<String>,
    pub guidance: Option<String>,
}

#[derive(Clone, Debug)]
pub struct RunOptions {
    pub max_turns: i64,
    pub concurrency: usize,
    /// Which agent is under test (`{name, version, model, ...}`). Pure
    /// telemetry, stored on every rollout; strongly recommended.
    pub agent_meta: Option<Value>,
}

impl Default for RunOptions {
    fn default() -> Self {
        Self {
            max_turns: 12,
            concurrency: 4,
            agent_meta: None,
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct ScenarioOptions {
    pub max_turns: Option<i64>,
    pub random_seed: Option<i64>,
    pub agent_meta: Option<Value>,
    pub dataset_id: Option<String>,
}

#[derive(Clone, Debug)]
pub struct PrepareOptions {
    pub count: i64,
    pub max_turns: i64,
    pub min_success_rate: f64,
    pub max_success_rate: f64,
    /// Force the full refresh: re-interview, re-distill, fresh batch.
    pub reprobe: bool,
    /// Part of the reuse key — a dataset probed under different grounding is
    /// not reused. `None` leaves reuse exactly as it was.
    pub seed_ids: Option<Vec<String>>,
    pub verbose: bool,
}

impl Default for PrepareOptions {
    fn default() -> Self {
        Self {
            count: 100,
            max_turns: 10,
            min_success_rate: 0.6,
            max_success_rate: 0.9,
            reprobe: false,
            seed_ids: None,
            verbose: false,
        }
    }
}

#[derive(Clone)]
pub struct EvalOptions<'a> {
    pub count: i64,
    /// Roll out this dataset instead of preparing one.
    pub dataset: Option<String>,
    /// Probe agent for user-model creation. Defaults to driving the rollout
    /// agent itself: each probe question becomes a one-turn conversation and the
    /// sandbox calls it makes are traced onto the reply. An audio agent MUST set
    /// this, since probing needs a text reply.
    pub probe_agent: Option<&'a dyn ProbeAgent>,
    pub reprobe: bool,
    pub seed_ids: Option<Vec<String>>,
    pub max_turns: i64,
    pub probe_max_turns: i64,
    pub concurrency: usize,
    pub repeats: usize,
    pub min_success_rate: f64,
    pub max_success_rate: f64,
    pub label: Option<String>,
    pub agent_meta: Option<Value>,
    pub verbose: bool,
}

impl std::fmt::Debug for EvalOptions<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // `probe_agent` is a trait object with no Debug bound.
        f.debug_struct("EvalOptions")
            .field("count", &self.count)
            .field("dataset", &self.dataset)
            .field("probe_agent", &self.probe_agent.map(|_| "<agent>"))
            .field("reprobe", &self.reprobe)
            .field("seed_ids", &self.seed_ids)
            .field("max_turns", &self.max_turns)
            .field("concurrency", &self.concurrency)
            .field("repeats", &self.repeats)
            .field("label", &self.label)
            .finish_non_exhaustive()
    }
}

impl Default for EvalOptions<'_> {
    fn default() -> Self {
        Self {
            count: 100,
            dataset: None,
            probe_agent: None,
            reprobe: false,
            seed_ids: None,
            max_turns: 12,
            probe_max_turns: 10,
            concurrency: 4,
            repeats: 1,
            min_success_rate: 0.6,
            max_success_rate: 0.9,
            label: None,
            agent_meta: None,
            verbose: false,
        }
    }
}

/// Outcome of [`Synthia::prepare`].
#[derive(Clone, Debug)]
pub struct PrepareResult {
    pub dataset: Dataset,
    pub user_model: types::UserModel,
    /// `"generated"` or `"reused"`.
    pub action: String,
    /// Human-readable decision trail — log it in CI.
    pub reason: String,
    pub success_rate: Option<f64>,
    pub quality_check_id: Option<String>,
}

/// Outcome of [`Synthia::run`]: everything each step produced, judged.
#[derive(Clone, Debug)]
pub struct EvalOutcome {
    /// `None` when a dataset was passed in.
    pub prepare: Option<PrepareResult>,
    pub dataset: Dataset,
    pub results: Vec<RolloutResult>,
    pub quality_check: QualityCheck,
    pub evaluations: Vec<types::RolloutEvaluation>,
    /// Judged pass fraction; `None` when nothing was judged.
    pub pass_rate: Option<f64>,
}

/// One finished rollout.
#[derive(Clone, Debug)]
pub struct RolloutResult {
    pub rollout_id: String,
    pub scenario_id: String,
    pub status: String,
    pub turns: i64,
    pub transcript: Vec<types::TranscriptTurn>,
    pub tool_events: Vec<types::ToolCall>,
    /// Voice mode: true when the rollout ran with voiced replies.
    pub voiced: bool,
    pub voice_render_id: Option<String>,
    http: Arc<Http>,
}

impl RolloutResult {
    /// Download the mixed conversation WAV, waiting for the server-attached
    /// render. `Err` on a rollout that isn't voiced — the peers express this as
    /// an absent property (JS) / a raise (Python).
    pub async fn audio(&self, opts: WaitOptions) -> Result<Vec<u8>> {
        let render_id = self.voice_render_id.as_deref().ok_or_else(|| {
            Error::Invalid(format!(
                "rollout {} is not voiced; there is no render to fetch",
                self.rollout_id
            ))
        })?;
        fetch_render_audio(&self.http, render_id, opts).await
    }

    /// Write the mixed conversation WAV to `path`; returns the path.
    ///
    /// Deliberately not generic over `AsRef<Path>`: the cross-SDK surface check
    /// greps for `fn save_audio(`, and a generic parameter list hides the method
    /// from it. Keeping the signature plain keeps the gate honest.
    pub async fn save_audio(&self, path: &Path, opts: WaitOptions) -> Result<PathBuf> {
        let bytes = self.audio(opts).await?;
        std::fs::write(path, bytes)?;
        Ok(path.to_path_buf())
    }
}

async fn fetch_render_audio(http: &Http, render_id: &str, opts: WaitOptions) -> Result<Vec<u8>> {
    let deadline = tokio::time::Instant::now() + Duration::from_secs_f64(opts.timeout);
    let mut render: types::VoiceRender = http
        .get(&format!("/v1/voice-renders/{render_id}"), &[])
        .await?;
    while render.status == types::VoiceRenderStatus::Running {
        if tokio::time::Instant::now() > deadline {
            return Err(Error::Timeout {
                kind: "voice render".to_string(),
                id: render_id.to_string(),
                secs: opts.timeout,
            });
        }
        tokio::time::sleep(Duration::from_secs_f64(opts.poll_interval)).await;
        render = http
            .get(&format!("/v1/voice-renders/{render_id}"), &[])
            .await?;
    }
    if render.status != types::VoiceRenderStatus::Succeeded {
        return Err(Error::Failed {
            kind: "voice render".to_string(),
            id: render_id.to_string(),
            error: render.error.clone(),
        });
    }
    http.get_bytes(&format!("/v1/voice-renders/{render_id}/audio"))
        .await
}

// --- jobs and datasets ------------------------------------------------------

/// A generated dataset of scenarios.
#[derive(Clone, Debug)]
pub struct Dataset {
    pub id: String,
    pub generation_id: String,
    pub user_model_id: String,
    pub row_count: i64,
    http: Arc<Http>,
}

impl Dataset {
    fn new(data: types::Dataset, http: Arc<Http>) -> Self {
        Self {
            id: data.id,
            generation_id: data.generation_id,
            user_model_id: data.user_model_id,
            row_count: data.row_count,
            http,
        }
    }

    pub async fn download(&self) -> Result<Vec<types::DatasetRow>> {
        let body: types::DatasetRows = self
            .http
            .get(&format!("/v1/datasets/{}/rows", self.id), &[])
            .await?;
        Ok(body.data)
    }

    /// Start an async validation run over this dataset. `label` is a human name
    /// shown on the platform's Runs page.
    pub async fn validate(&self, label: Option<&str>) -> Result<ValidationRun> {
        // Two explicit call sites rather than one with a computed body: the
        // peers send NO body when there is no label (JS `undefined`, Python
        // `None`), and scripts/check-request-coverage.py reads the keys
        // lexically beside the path.
        let path = format!("/v1/datasets/{}/validations", self.id);
        let data: types::ValidationRun = match label {
            Some(l) => self.http.post(&path, json!({ "label": l })).await?,
            None => self.http.post(&path, Value::Null).await?,
        };
        Ok(ValidationRun::new(data, self.http.clone()))
    }

    /// Play this dataset's scenarios against a local agent.
    pub async fn rollout(
        &self,
        agent: &dyn RolloutAgent,
        opts: RunOptions,
    ) -> Result<Vec<RolloutResult>> {
        Rollouts {
            http: self.http.clone(),
            session_id: None,
        }
        .run(agent, Some(&self.id), opts)
        .await
    }
}

/// An async scenario-generation job.
#[derive(Clone, Debug)]
pub struct GenerationJob {
    pub id: String,
    pub status: String,
    pub user_model_id: String,
    pub count: i64,
    pub dataset_id: Option<String>,
    pub error: Option<String>,
    http: Arc<Http>,
}

impl GenerationJob {
    fn new(data: types::GenerationJob, http: Arc<Http>) -> Self {
        Self {
            id: data.id,
            status: status_str(&data.status),
            user_model_id: data.user_model_id,
            count: data.count,
            dataset_id: data.dataset_id,
            error: data.error,
            http,
        }
    }

    /// Poll until the job finishes; `verbose` prints server telemetry live.
    pub async fn wait(mut self, opts: WaitOptions) -> Result<Dataset> {
        let mut events = EventStream::new(
            self.http.clone(),
            format!("/v1/generations/{}/events", self.id),
            opts.verbose,
        );
        let deadline = tokio::time::Instant::now() + Duration::from_secs_f64(opts.timeout);
        while self.status == "running" {
            if tokio::time::Instant::now() > deadline {
                return Err(Error::Timeout {
                    kind: "generation".to_string(),
                    id: self.id,
                    secs: opts.timeout,
                });
            }
            tokio::time::sleep(Duration::from_secs_f64(opts.poll_interval)).await;
            let data: types::GenerationJob = self
                .http
                .get(&format!("/v1/generations/{}", self.id), &[])
                .await?;
            self.status = status_str(&data.status);
            self.dataset_id = data.dataset_id;
            self.error = data.error;
            events.pump().await?;
        }
        events.pump().await?; // catch events written after the final status poll
        if self.status != "succeeded" {
            return Err(Error::Failed {
                kind: "generation".to_string(),
                id: self.id,
                error: self.error,
            });
        }
        let dataset_id = self
            .dataset_id
            .ok_or_else(|| Error::Api("generation succeeded without a dataset".to_string()))?;
        let data: types::Dataset = self
            .http
            .get(&format!("/v1/datasets/{dataset_id}"), &[])
            .await?;
        Ok(Dataset::new(data, self.http))
    }
}

/// A validation of a dataset: a per-scenario judge gate plus collective
/// fidelity/diversity reports and an advisory verdict.
#[derive(Clone, Debug)]
pub struct ValidationRun {
    pub id: String,
    pub dataset_id: String,
    pub status: String,
    pub label: Option<String>,
    pub verdict: Option<String>,
    pub reference: Option<Map<String, Value>>,
    pub validity: Option<Map<String, Value>>,
    pub fidelity: Option<Map<String, Value>>,
    pub diversity: Option<Map<String, Value>>,
    pub error: Option<String>,
    http: Arc<Http>,
}

impl ValidationRun {
    fn new(data: types::ValidationRun, http: Arc<Http>) -> Self {
        Self {
            id: data.id,
            dataset_id: data.dataset_id,
            status: status_str(&data.status),
            label: data.label,
            verdict: data.verdict.as_ref().map(status_str),
            reference: data.reference,
            validity: data.validity,
            fidelity: data.fidelity,
            diversity: data.diversity,
            error: data.error,
            http,
        }
    }

    pub async fn wait(mut self, opts: WaitOptions) -> Result<ValidationRun> {
        let mut events = EventStream::new(
            self.http.clone(),
            format!("/v1/validation-runs/{}/events", self.id),
            opts.verbose,
        );
        let deadline = tokio::time::Instant::now() + Duration::from_secs_f64(opts.timeout);
        while self.status == "running" {
            if tokio::time::Instant::now() > deadline {
                return Err(Error::Timeout {
                    kind: "validation run".to_string(),
                    id: self.id,
                    secs: opts.timeout,
                });
            }
            tokio::time::sleep(Duration::from_secs_f64(opts.poll_interval)).await;
            let data: types::ValidationRun = self
                .http
                .get(&format!("/v1/validation-runs/{}", self.id), &[])
                .await?;
            self.status = status_str(&data.status);
            self.verdict = data.verdict.as_ref().map(status_str);
            self.reference = data.reference;
            self.validity = data.validity;
            self.fidelity = data.fidelity;
            self.diversity = data.diversity;
            self.error = data.error;
            events.pump().await?;
        }
        events.pump().await?;
        if self.status != "succeeded" {
            return Err(Error::Failed {
                kind: "validation run".to_string(),
                id: self.id,
                error: self.error,
            });
        }
        Ok(self)
    }

    /// Per-scenario judge verdicts: scenario_id, passed, judge.
    pub async fn scenarios(&self) -> Result<Vec<types::ScenarioValidation>> {
        let body: types::ScenarioValidationList = self
            .http
            .get(&format!("/v1/validation-runs/{}/scenarios", self.id), &[])
            .await?;
        Ok(body.data)
    }
}

/// An async evaluation of finished rollouts. The per-rollout results are the
/// product; there is no aggregate verdict.
#[derive(Clone, Debug)]
pub struct QualityCheck {
    pub id: String,
    pub status: String,
    pub rollout_ids: Vec<String>,
    pub label: Option<String>,
    pub error: Option<String>,
    http: Arc<Http>,
}

impl QualityCheck {
    fn new(data: types::QualityCheck, http: Arc<Http>) -> Self {
        Self {
            id: data.id,
            status: status_str(&data.status),
            rollout_ids: data.rollout_ids,
            label: data.label,
            error: data.error,
            http,
        }
    }

    pub async fn wait(&mut self, opts: WaitOptions) -> Result<()> {
        let mut events = EventStream::new(
            self.http.clone(),
            format!("/v1/quality-checks/{}/events", self.id),
            opts.verbose,
        );
        let deadline = tokio::time::Instant::now() + Duration::from_secs_f64(opts.timeout);
        while self.status == "running" {
            if tokio::time::Instant::now() > deadline {
                return Err(Error::Timeout {
                    kind: "quality check".to_string(),
                    id: self.id.clone(),
                    secs: opts.timeout,
                });
            }
            tokio::time::sleep(Duration::from_secs_f64(opts.poll_interval)).await;
            let data: types::QualityCheck = self
                .http
                .get(&format!("/v1/quality-checks/{}", self.id), &[])
                .await?;
            self.status = status_str(&data.status);
            self.error = data.error;
            events.pump().await?;
        }
        events.pump().await?;
        if self.status != "succeeded" {
            return Err(Error::Failed {
                kind: "quality check".to_string(),
                id: self.id.clone(),
                error: self.error.clone(),
            });
        }
        Ok(())
    }

    /// Per-rollout results: rollout_id, passed, states (the agentic-state
    /// trajectory), and judge (dimensions + issues).
    pub async fn rollouts(&self) -> Result<Vec<types::RolloutEvaluation>> {
        let body: types::RolloutEvaluationList = self
            .http
            .get(&format!("/v1/quality-checks/{}/rollouts", self.id), &[])
            .await?;
        Ok(body.data)
    }
}

/// Open enums render as their wire string; `Other` keeps an unknown value
/// verbatim, so a newer server never breaks a status comparison.
fn status_str<T: serde::Serialize>(status: &T) -> String {
    match serde_json::to_value(status) {
        Ok(Value::String(s)) => s,
        Ok(other) => other.to_string(),
        Err(_) => String::new(),
    }
}

// --- resources --------------------------------------------------------------

pub struct Seeds {
    http: Arc<Http>,
}

impl Seeds {
    /// Upload seed material (documents, tool schemas, policies, traces...).
    ///
    /// JSON content is ingested as-is; bytes or a path are uploaded and handled
    /// server-side with zero parameters — audio is transcribed (and marks the
    /// seed voice-origin), everything else is rendered to text.
    pub async fn ingest(
        &self,
        kind: &str,
        source: &str,
        content: Content,
        version: Option<&str>,
        metadata: Option<Value>,
    ) -> Result<types::Seed> {
        let version = version.unwrap_or("1");
        let metadata = metadata.unwrap_or_else(|| json!({}));
        let upload = match &content {
            Content::Json(_) => None,
            Content::Bytes(b) => as_file_upload(&Reply::Bytes(b.clone()))?,
            Content::File(p) => {
                let up = as_file_upload(&Reply::File(p.clone()))?;
                if up.is_none() {
                    return Err(Error::Invalid(format!(
                        "{} is not a supported upload type",
                        p.display()
                    )));
                }
                up
            }
        };
        if let Some(up) = upload {
            return self
                .http
                .post(
                    "/v1/seeds",
                    json!({
                        "kind": kind,
                        "source": source,
                        "file_b64": up.file_b64,
                        "filename": up.filename,
                        "version": version,
                        "metadata": metadata,
                    }),
                )
                .await;
        }
        let Content::Json(value) = content else {
            unreachable!("non-JSON content produced no upload")
        };
        self.http
            .post(
                "/v1/seeds",
                json!({
                    "kind": kind,
                    "source": source,
                    "content": value,
                    "version": version,
                    "metadata": metadata,
                }),
            )
            .await
    }
}

pub struct UserModels {
    http: Arc<Http>,
}

impl UserModels {
    /// Probe `agent` until the server converges on a user model.
    ///
    /// The agent runs locally; only probe questions, replies, and traced tool
    /// calls travel over the wire.
    pub async fn create_from_probe(
        &self,
        agent: &dyn ProbeAgent,
        opts: ProbeOptions,
    ) -> Result<types::UserModel> {
        let mut body = json!({ "max_turns": opts.max_turns });
        if let Some(ids) = &opts.seed_ids {
            body["seed_ids"] = json!(ids);
        }
        let mut session: types::ProbeSession = self.http.post("/v1/probe-sessions", body).await?;
        if let Some(requested) = &opts.seed_ids {
            if !same_id_set(&session.seed_ids, requested) {
                // An older deployment that drops the field would otherwise leave
                // every run silently ungrounded — the exact failure this option
                // exists to fix.
                return Err(Error::Api(format!(
                    "server did not accept seed_ids (echoed {:?}, requested {:?}); \
                     it is likely older than the grounding feature",
                    session.seed_ids, requested
                )));
            }
        }
        let mut events = EventStream::new(
            self.http.clone(),
            format!("/v1/probe-sessions/{}/events", session.id),
            opts.verbose,
        );
        while session.status == types::ProbeSessionStatus::Active {
            let probe = session.next_probe.clone().ok_or_else(|| {
                Error::Api("probe session is active with no next_probe".to_string())
            })?;
            let answer = agent.answer(probe).await?;
            session = self
                .http
                .post(
                    &format!("/v1/probe-sessions/{}/responses", session.id),
                    json!({ "reply": answer.reply, "tool_calls": answer.tool_calls }),
                )
                .await?;
            events.pump().await?;
        }
        let model_id = session
            .user_model_id
            .ok_or_else(|| Error::Api("probe converged without a user model".to_string()))?;
        self.get(&model_id).await
    }

    pub async fn get(&self, model_id: &str) -> Result<types::UserModel> {
        self.http
            .get(&format!("/v1/user-models/{model_id}"), &[])
            .await
    }

    /// User models, **oldest first** — the newest is the last element.
    pub async fn list(&self, session: Option<&str>) -> Result<Vec<types::UserModel>> {
        let params: Vec<(&str, String)> = match session {
            Some(s) => vec![("sdk_session", s.to_string())],
            None => vec![],
        };
        let body: types::UserModelList = self.http.get("/v1/user-models", &params).await?;
        Ok(body.data)
    }

    /// Submit user-authored scenarios against a user model. The returned
    /// `passed`/`judge` are ADVISORY — a failing verdict never blocks.
    pub async fn submit_scenarios(
        &self,
        user_model: impl ModelRef,
        scenarios: Vec<Value>,
    ) -> Result<Vec<types::CustomScenarioResult>> {
        let body: types::CustomScenariosResult = self
            .http
            .post(
                &format!("/v1/user-models/{}/scenarios", user_model.model_id()),
                json!({ "scenarios": scenarios }),
            )
            .await?;
        Ok(body.data)
    }
}

pub struct Datasets {
    http: Arc<Http>,
}

impl Datasets {
    pub async fn get(&self, dataset_id: &str) -> Result<Dataset> {
        let data: types::Dataset = self
            .http
            .get(&format!("/v1/datasets/{dataset_id}"), &[])
            .await?;
        Ok(Dataset::new(data, self.http.clone()))
    }

    /// Datasets **newest first**; `session` filters to one SDK session.
    pub async fn list(&self, session: Option<&str>) -> Result<Vec<Dataset>> {
        let params: Vec<(&str, String)> = match session {
            Some(s) => vec![("sdk_session", s.to_string())],
            None => vec![],
        };
        let body: types::DatasetList = self.http.get("/v1/datasets", &params).await?;
        Ok(body
            .data
            .into_iter()
            .map(|d| Dataset::new(d, self.http.clone()))
            .collect())
    }

    /// Start a generation job. `quality_check_id` names a completed check whose
    /// results calibrate the batch's difficulty and coverage. `guidance` biases
    /// scenario content toward a theme; it never overrides grounding.
    pub async fn generate(
        &self,
        user_model: impl ModelRef,
        opts: GenerateOptions,
    ) -> Result<GenerationJob> {
        let mut body = json!({
            "user_model_id": user_model.model_id(),
            "count": opts.count.unwrap_or(20),
        });
        if let Some(qc) = &opts.quality_check_id {
            body["quality_check_id"] = json!(qc);
        }
        if let Some(g) = &opts.guidance {
            body["guidance"] = json!(g);
        }
        let data: types::GenerationJob = self.http.post("/v1/generations", body).await?;
        Ok(GenerationJob::new(data, self.http.clone()))
    }

    /// Assemble a dataset from an explicit set of scenarios that all belong to
    /// `user_model`'s representation.
    pub async fn compose(
        &self,
        user_model: impl ModelRef,
        scenario_ids: Vec<String>,
        label: Option<&str>,
    ) -> Result<Dataset> {
        let mut body = json!({
            "user_model_id": user_model.model_id(),
            "scenario_ids": scenario_ids,
        });
        if let Some(l) = label {
            body["label"] = json!(l);
        }
        let data: types::Dataset = self.http.post("/v1/datasets", body).await?;
        Ok(Dataset::new(data, self.http.clone()))
    }
}

pub struct Rollouts {
    http: Arc<Http>,
    session_id: Option<String>,
}

impl Rollouts {
    /// A stored rollout's full captured state.
    pub async fn get(&self, rollout_id: &str) -> Result<types::Rollout> {
        self.http
            .get(&format!("/v1/rollouts/{rollout_id}"), &[])
            .await
    }

    /// One turn's audio (turns with an `audio_url` only).
    pub async fn turn_audio(&self, rollout_id: &str, idx: i64) -> Result<Vec<u8>> {
        self.http
            .get_bytes(&format!("/v1/rollouts/{rollout_id}/turns/{idx}/audio"))
            .await
    }

    /// Play a dataset's scenarios against `agent` (most recent dataset when
    /// none is given).
    ///
    /// Scenarios run with `concurrency` in flight at once; turns within one
    /// conversation are sequential. Results keep dataset row order regardless of
    /// completion order.
    pub async fn run(
        &self,
        agent: &dyn RolloutAgent,
        dataset: Option<&str>,
        opts: RunOptions,
    ) -> Result<Vec<RolloutResult>> {
        let dataset_id = match dataset {
            Some(id) => id.to_string(),
            None => {
                // Session-scoped default: this script's latest dataset, so two
                // concurrent scripts never pick up each other's data.
                let mut data: Vec<types::Dataset> = Vec::new();
                if let Some(sid) = &self.session_id {
                    let body: types::DatasetList = self
                        .http
                        .get("/v1/datasets", &[("sdk_session", sid.clone())])
                        .await?;
                    data = body.data;
                }
                if data.is_empty() {
                    let body: types::DatasetList = self.http.get("/v1/datasets", &[]).await?;
                    data = body.data;
                    if !data.is_empty() && self.session_id.is_some() {
                        println!(
                            "note: no dataset in this session yet; using latest dataset {}",
                            data[0].id
                        );
                    }
                }
                if data.is_empty() {
                    return Err(Error::Api(
                        "no datasets exist yet; generate one first".to_string(),
                    ));
                }
                data[0].id.clone()
            }
        };
        let body: types::DatasetRows = self
            .http
            .get(&format!("/v1/datasets/{dataset_id}/rows"), &[])
            .await?;
        let rows = body.data;

        use futures::stream::StreamExt;
        let concurrency = opts.concurrency.max(1).min(rows.len().max(1));
        let mut indexed: Vec<(usize, RolloutResult)> =
            futures::stream::iter(rows.iter().enumerate().map(|(i, row)| {
                let opts = opts.clone();
                let dataset_id = dataset_id.clone();
                async move {
                    let result = self
                        .run_scenario(
                            agent,
                            &row.scenario_id,
                            ScenarioOptions {
                                max_turns: Some(opts.max_turns),
                                random_seed: row.random_seed,
                                agent_meta: opts.agent_meta.clone(),
                                dataset_id: Some(dataset_id),
                            },
                        )
                        .await?;
                    Ok::<(usize, RolloutResult), Error>((i, result))
                }
            }))
            .buffer_unordered(concurrency)
            .collect::<Vec<_>>()
            .await
            .into_iter()
            .collect::<Result<Vec<_>>>()?;

        indexed.sort_by_key(|(i, _)| *i);
        Ok(indexed.into_iter().map(|(_, r)| r).collect())
    }

    /// Start an async quality check over finished rollouts.
    pub async fn quality_check<R: RolloutRef>(
        &self,
        rollouts: &[R],
        label: Option<&str>,
    ) -> Result<QualityCheck> {
        let ids: Vec<&str> = rollouts.iter().map(RolloutRef::rollout_id).collect();
        let data: types::QualityCheck = self
            .http
            .post(
                "/v1/quality-checks",
                json!({ "rollout_ids": ids, "label": label }),
            )
            .await?;
        Ok(QualityCheck::new(data, self.http.clone()))
    }

    /// Recover a turn post that failed transiently, without risking a
    /// double-advance. Re-fetch the rollout: if its turn count moved past
    /// `prior_turn`, the write landed (only the response was lost) — adopt that
    /// state. If it's unchanged and still running, it's safe to re-send once. A
    /// genuine 4xx (409 already-completed, 422 bad file/turn) is returned
    /// unchanged.
    async fn recover_turn(
        &self,
        rollout_id: &str,
        prior_turn: i64,
        body: &Value,
        original: Error,
    ) -> Result<types::Rollout> {
        if !original.is_transient() {
            return Err(original);
        }
        let fetched: types::Rollout = match self
            .http
            .get(&format!("/v1/rollouts/{rollout_id}"), &[])
            .await
        {
            Ok(r) => r,
            // Can't confirm state — surface the original failure.
            Err(_) => return Err(original),
        };
        if fetched.status != types::RolloutStatus::Running || fetched.turn > prior_turn {
            return Ok(fetched); // the turn landed (or the rollout finished)
        }
        // The turn never reached the server — safe to send it again.
        self.http
            .post(&format!("/v1/rollouts/{rollout_id}/turns"), body.clone())
            .await
    }

    /// Run one rollout session; one HTTP round-trip per agent turn.
    pub async fn run_scenario(
        &self,
        agent: &dyn RolloutAgent,
        scenario_id: &str,
        opts: ScenarioOptions,
    ) -> Result<RolloutResult> {
        let mut session: types::Rollout = self
            .http
            .post(
                "/v1/rollouts",
                json!({
                    "scenario_id": scenario_id,
                    "random_seed": opts.random_seed,
                    "max_turns": opts.max_turns.unwrap_or(12),
                    "agent": opts.agent_meta,
                    "dataset_id": opts.dataset_id,
                }),
            )
            .await?;

        // ONE sandbox for the whole rollout: the agent's per-conversation state
        // lives on sandbox.context, so identity must be stable across turns (a
        // fresh object per turn silently amnesia-ed every stateful agent).
        // Server-mirrored fields refresh each turn.
        let mut sandbox: Option<ToolSandbox> = None;
        while session.status == types::RolloutStatus::Running {
            match sandbox.as_mut() {
                None => sandbox = Some(ToolSandbox::from_config(&session.sandbox)),
                Some(sb) => {
                    sb.state = session.sandbox.state.clone();
                    sb.fail_tools = session.sandbox.fail_tools.iter().cloned().collect();
                    sb.events.clear();
                }
            }
            let sb = sandbox.as_mut().expect("sandbox initialized above");
            let reply = agent.reply(session.transcript.clone(), sb).await?;
            let upload = as_file_upload(&reply)?;
            let mut body = json!({ "tool_calls": sb.events });
            match upload {
                Some(up) => {
                    // The server sniffs the file: audio is transcribed and flips
                    // the rollout into voice mode; images and documents are
                    // rendered to text.
                    body["reply"] = json!("");
                    body["file_b64"] = json!(up.file_b64);
                    body["filename"] = json!(up.filename);
                }
                None => {
                    // `Reply::File` is unambiguous intent, unlike the peers'
                    // `string`, which doubles as both text and a path — there,
                    // an unrecognized suffix silently becomes the reply TEXT
                    // (the path itself). Rust can tell the two apart, so refuse
                    // rather than send a path, or an empty reply, as the answer.
                    body["reply"] = match &reply {
                        Reply::Text(t) => json!(t),
                        Reply::File(p) => {
                            return Err(Error::Invalid(format!(
                                "{} is not a supported reply file type ({}); \
                                 return Reply::Text for a text reply",
                                p.display(),
                                FILE_SUFFIXES.join(", ")
                            )))
                        }
                        Reply::Bytes(_) => unreachable!("bytes always upload"),
                    };
                }
            }
            // Turn posts advance server transcript state, so a blind retry could
            // double-record the reply — the plain retry excludes them. Recover
            // explicitly instead (check-then-act idempotency).
            let prior_turn = session.turn;
            let posted: Result<types::Rollout> = self
                .http
                .post(&format!("/v1/rollouts/{}/turns", session.id), body.clone())
                .await;
            session = match posted {
                Ok(s) => s,
                Err(e) => self.recover_turn(&session.id, prior_turn, &body, e).await?,
            };
        }

        Ok(RolloutResult {
            rollout_id: session.id,
            scenario_id: scenario_id.to_string(),
            status: status_str(&session.status),
            turns: session.turn,
            transcript: session.transcript,
            tool_events: session.tool_events,
            voiced: session.voiced,
            voice_render_id: session.voice_render_id,
            http: self.http.clone(),
        })
    }
}

// --- production issue traces ------------------------------------------------

/// Detached trace recorder with `report()` ergonomics identical to
/// [`ToolSandbox`], but not bound to any rollout.
#[derive(Clone, Debug, Default)]
pub struct TraceRecorder {
    turns: Vec<types::TranscriptTurn>,
    events: Vec<Value>,
}

impl TraceRecorder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn user(&mut self, content: &str) {
        self.turns.push(types::TranscriptTurn {
            role: types::TranscriptTurnRole::User,
            content: content.to_string(),
            audio_url: None,
        });
    }

    pub fn agent(&mut self, content: &str) {
        self.turns.push(types::TranscriptTurn {
            role: types::TranscriptTurnRole::Agent,
            content: content.to_string(),
            audio_url: None,
        });
    }

    /// Record a real tool call the agent made on the current turn.
    pub fn report(&mut self, name: &str, output: Map<String, Value>, opts: crate::Report) {
        self.events.push(json!({
            "name": name,
            "input": opts.input.unwrap_or_default(),
            "output": output,
            "is_error": opts.is_error,
            "external": true,
            "turn_idx": self.turns.len().saturating_sub(1),
        }));
    }
}

/// Options for [`Traces::ingest`]. Mirrors the peers' options object.
#[derive(Default)]
pub struct TraceIngestOptions<'a> {
    /// Provenance label shown on the Issues page (e.g. "prod/zendesk").
    pub source: &'a str,
    /// Optional operator hint describing the failure; sharpens extraction.
    pub error: Option<&'a str>,
    /// Size of the generated dataset (default 20).
    pub count: Option<i64>,
    /// Extra redactor composed AFTER the built-in detectors, applied to every
    /// transcript/tool string before the trace leaves this process. Names and
    /// addresses need this — the structured detectors cannot catch them.
    pub redact: Option<&'a (dyn Fn(&str) -> String + Send + Sync)>,
}

pub struct Traces {
    http: Arc<Http>,
}

impl Traces {
    /// Ingest a production issue trace: redact it client-side (required, always
    /// on), then distill it server-side into synthetic scenarios that reproduce
    /// the failure. The raw trace is never persisted.
    pub async fn ingest(
        &self,
        trace: &TraceRecorder,
        opts: TraceIngestOptions<'_>,
    ) -> Result<GenerationJob> {
        let extra = opts.redact;
        let transcript: Vec<Value> = trace
            .turns
            .iter()
            .map(|t| {
                json!({
                    "role": role_str(&t.role),
                    "content": crate::redact::redact_string_with(&t.content, extra),
                })
            })
            .collect();
        let tool_events: Vec<Value> = trace
            .events
            .iter()
            .map(|e| crate::redact::redact_value_with(e, extra))
            .collect();
        let data: types::GenerationJob = self
            .http
            .post(
                "/v1/traces",
                json!({
                    "transcript": transcript,
                    "tool_events": tool_events,
                    "source": opts.source,
                    "error": opts.error,
                    "count": opts.count.unwrap_or(20),
                }),
            )
            .await?;
        Ok(GenerationJob::new(data, self.http.clone()))
    }
}

// --- client -----------------------------------------------------------------

/// Session name derived from the entry point: `"<parent-dir>/<binary-stem>"`,
/// falling back to the current directory's name. Mirrors the JS
/// `defaultSessionName()` and Python `_default_session_name()`.
fn default_session_name() -> String {
    let exe = std::env::current_exe().ok();
    let stem = exe
        .as_ref()
        .and_then(|p| p.file_stem())
        .and_then(|s| s.to_str())
        .unwrap_or("");
    let cwd_name = std::env::current_dir()
        .ok()
        .and_then(|p| p.file_name().and_then(|n| n.to_str()).map(str::to_string))
        .unwrap_or_else(|| "synthia".to_string());
    // A cargo test/example binary stem is noise, not an identity.
    if stem.is_empty() || stem.starts_with("deps") {
        return cwd_name;
    }
    let parent = exe
        .as_ref()
        .and_then(|p| p.parent())
        .and_then(|p| p.file_name())
        .and_then(|n| n.to_str())
        .unwrap_or("");
    if parent.is_empty() || parent == "debug" || parent == "release" || parent == "deps" {
        format!("{cwd_name}/{stem}")
    } else {
        format!("{parent}/{stem}")
    }
}

/// Client entry point.
///
/// Session identity: every client belongs to a named session — the stable,
/// account-scoped identity of one script, persisted across executions (same name
/// resumes the same session; re-runs reuse its datasets instead of
/// re-probing/re-generating). Resolution order: the `session` option >
/// `SYNTHIA_SESSION` > a derived "project/script" name.
///
/// Degradation: an old server without `/v1/sdk-sessions` → no session; keyless
/// against a keyed server → anonymous session; an invalid key fails on
/// [`Synthia::ready`].
pub struct Synthia {
    pub session_name: String,
    pub session_id: Option<String>,
    pub invocation_id: Option<String>,
    /// CI floors/caps mirrored from the customer config by the handshake.
    pub ci_settings: Option<types::CiSettings>,
    pub seeds: Seeds,
    pub user_models: UserModels,
    pub datasets: Datasets,
    pub rollouts: Rollouts,
    pub traces: Traces,
    http: Arc<Http>,
    ci: Option<Value>,
    handshake_done: bool,
}

impl Synthia {
    pub fn new(options: SynthiaOptions) -> Result<Self> {
        let api_key = options
            .api_key
            .or_else(|| std::env::var("SYNTHIA_API_KEY").ok());
        let base_url = options
            .base_url
            .or_else(|| std::env::var("SYNTHIA_BASE_URL").ok())
            .unwrap_or_else(|| DEFAULT_BASE_URL.to_string());
        let http = Arc::new(Http::new(base_url, api_key)?);

        let session_name = if options.ephemeral_session {
            let suffix: String = {
                let mut rng = rand::thread_rng();
                (0..8)
                    .map(|_| format!("{:x}", rng.gen_range(0..16)))
                    .collect()
            };
            format!("{}/eph-{}", default_session_name(), suffix)
        } else if let Some(name) = options.session {
            name
        } else {
            std::env::var("SYNTHIA_SESSION")
                .ok()
                .filter(|s| !s.is_empty())
                .unwrap_or_else(default_session_name)
        };

        Ok(Self {
            session_name,
            session_id: None,
            invocation_id: None,
            ci_settings: None,
            seeds: Seeds { http: http.clone() },
            user_models: UserModels { http: http.clone() },
            datasets: Datasets { http: http.clone() },
            rollouts: Rollouts {
                http: http.clone(),
                session_id: None,
            },
            traces: Traces { http: http.clone() },
            http,
            ci: options.ci,
            handshake_done: false,
        })
    }

    /// Run the session handshake.
    ///
    /// Unlike Python — which handshakes synchronously in `__init__` — and like
    /// JS, construction does no I/O. Call this to fail fast on a bad key and to
    /// read `session_id` / `ci_settings` before acting. Every other method calls
    /// it implicitly if you don't.
    pub async fn ready(&mut self) -> Result<()> {
        if self.handshake_done {
            return Ok(());
        }
        self.handshake_done = true;
        // The serverless backend can 500 on the first request after idling, and
        // the handshake is every process's first request — retry briefly (which
        // also warms the container), then degrade to sessionless rather than
        // failing over optional tracking.
        let mut body = json!({
            "name": self.session_name,
            "sdk_version": SDK_VERSION,
        });
        if let Some(ci) = &self.ci {
            body["ci"] = ci.clone();
        }
        let mut last: Option<reqwest::Response> = None;
        for attempt in 0..3u32 {
            match self
                .http
                .raw("POST", "/v1/sdk-sessions", Some(&body), &[])
                .await
            {
                Ok(r) => {
                    let status = r.status().as_u16();
                    last = Some(r);
                    if status < 500 {
                        break;
                    }
                }
                Err(_) => last = None,
            }
            if attempt < 2 {
                tokio::time::sleep(Duration::from_secs(2 * (attempt as u64 + 1))).await;
            }
        }
        let Some(resp) = last else {
            return Ok(()); // transient trouble even after retries: run untracked
        };
        let status = resp.status().as_u16();
        if status >= 500 || status == 404 {
            // 404: server predates sessions. 5xx: degrade rather than fail.
            return Ok(());
        }
        if status == 401 {
            let text = resp.text().await.unwrap_or_default();
            let detail = serde_json::from_str::<Value>(&text)
                .ok()
                .and_then(|v| v["detail"].as_str().map(str::to_string))
                .unwrap_or_else(|| "invalid API key".to_string());
            return Err(Error::Api(detail));
        }
        let resp = Http::check(resp, "POST", "/v1/sdk-sessions").await?;
        let data: types::SdkSession = serde_json::from_str(&resp.text().await?)?;
        self.session_id = Some(data.sdk_session_id.clone());
        self.invocation_id = Some(data.sdk_invocation_id.clone());
        self.ci_settings = data.ci;
        self.http
            .set_header("X-Synthia-Session", &data.sdk_session_id)
            .await;
        self.http
            .set_header("X-Synthia-Invocation", &data.sdk_invocation_id)
            .await;
        self.rollouts = Rollouts {
            http: self.http.clone(),
            session_id: Some(data.sdk_session_id),
        };
        Ok(())
    }

    /// Probe + generate only when needed; otherwise reuse the latest dataset.
    ///
    /// `count` is exact: reuse requires the latest dataset to match it as well as
    /// passing the quality gate. All lookups are scoped to this client's session.
    pub async fn prepare(
        &mut self,
        agent: &dyn ProbeAgent,
        opts: PrepareOptions,
    ) -> Result<PrepareResult> {
        self.ready().await?;
        if opts.reprobe {
            // Explicit refresh: skip every reuse check.
            return self
                .probe_and_generate(
                    agent,
                    &opts,
                    "reprobe requested".to_string(),
                    None,
                    None,
                    true,
                )
                .await;
        }
        let existing = self.datasets.list(self.session_id.as_deref()).await?; // newest first
        if existing.is_empty() {
            let reason = if self.session_id.is_some() {
                "no datasets in this session yet"
            } else {
                "no datasets exist yet"
            };
            return self
                .probe_and_generate(agent, &opts, reason.to_string(), None, None, false)
                .await;
        }

        let params: Vec<(&str, String)> = match &self.session_id {
            Some(s) => vec![("sdk_session", s.clone())],
            None => vec![],
        };
        let latest: types::QualityCheckSummary =
            self.http.get("/v1/quality-checks/latest", &params).await?;
        let rate = match (&latest.id, latest.total) {
            (Some(_), total) if total > 0 => Some(latest.passed as f64 / total as f64),
            _ => None,
        };

        if let Some(r) = rate {
            if !(opts.min_success_rate <= r && r <= opts.max_success_rate) {
                let (direction, bound) = if r < opts.min_success_rate {
                    ("below", opts.min_success_rate)
                } else {
                    ("above", opts.max_success_rate)
                };
                let reason = format!(
                    "success rate {} {} {}; regenerating calibrated on {}",
                    pct(r),
                    direction,
                    pct(bound),
                    latest.id.clone().unwrap_or_default()
                );
                return self
                    .probe_and_generate(agent, &opts, reason, rate, latest.id.clone(), false)
                    .await;
            }
        }

        // Quality in band (or unjudged): reuse only on an exact size match;
        // otherwise regenerate at the requested count — without re-probing,
        // since nothing suggests the agent changed.
        if existing[0].row_count != opts.count {
            let reason = format!(
                "latest dataset has {} rows; requested {}",
                existing[0].row_count, opts.count
            );
            return self
                .probe_and_generate(agent, &opts, reason, rate, None, false)
                .await;
        }

        // A dataset generated under different grounding answers a different
        // question, so grounding is part of the reuse key too.
        let reuse_model = self.user_models.get(&existing[0].user_model_id).await?;
        if !self
            .grounding_matches(&reuse_model.probe_session_id, opts.seed_ids.as_deref())
            .await?
        {
            return self
                .probe_and_generate(
                    agent,
                    &opts,
                    "requested grounding differs from the latest dataset's; \
                     re-probing and regenerating"
                        .to_string(),
                    rate,
                    None,
                    true,
                )
                .await;
        }

        let reason = match rate {
            Some(r) => format!(
                "success rate {} within {}-{} band",
                pct(r),
                pct(opts.min_success_rate),
                pct(opts.max_success_rate)
            ),
            None => "no completed quality check to judge by; reusing latest dataset".to_string(),
        };
        Ok(PrepareResult {
            dataset: existing[0].clone(),
            user_model: reuse_model,
            action: "reused".to_string(),
            reason,
            success_rate: rate,
            quality_check_id: None,
        })
    }

    /// The whole evaluation in one call: prepare → roll out every scenario →
    /// judge the rollouts → return the judged results. The script-path
    /// equivalent of `synthia run`, minus the CI gating.
    pub async fn run(
        &mut self,
        agent: &dyn RolloutAgent,
        opts: EvalOptions<'_>,
    ) -> Result<EvalOutcome> {
        if opts.dataset.is_some() && opts.seed_ids.is_some() {
            return Err(Error::Invalid(
                "seed_ids has no effect with dataset — passing a dataset skips \
                 prepare(), and grounding is applied while probing"
                    .to_string(),
            ));
        }
        self.ready().await?;
        let mut prepare: Option<PrepareResult> = None;
        let target = match &opts.dataset {
            Some(id) => self.datasets.get(id).await?,
            None => {
                // Default: drive the rollout agent itself as the prober.
                let fallback = ProbeFromRollout(agent);
                let probe: &dyn ProbeAgent = opts.probe_agent.unwrap_or(&fallback);
                let result = self
                    .prepare(
                        probe,
                        PrepareOptions {
                            count: opts.count,
                            max_turns: opts.probe_max_turns,
                            min_success_rate: opts.min_success_rate,
                            max_success_rate: opts.max_success_rate,
                            reprobe: opts.reprobe,
                            seed_ids: opts.seed_ids.clone(),
                            verbose: opts.verbose,
                        },
                    )
                    .await?;
                let dataset = result.dataset.clone();
                prepare = Some(result);
                dataset
            }
        };

        let mut results: Vec<RolloutResult> = Vec::new();
        for _ in 0..opts.repeats.max(1) {
            results.extend(
                self.rollouts
                    .run(
                        agent,
                        Some(&target.id),
                        RunOptions {
                            max_turns: opts.max_turns,
                            concurrency: opts.concurrency,
                            agent_meta: opts.agent_meta.clone(),
                        },
                    )
                    .await?,
            );
        }

        // The server bounds one quality check's LLM fan-out; bigger runs judge
        // in chunks and pool the evaluations. The outcome carries the last
        // check; every chunk lands on the platform.
        let chunks: Vec<&[RolloutResult]> = results.chunks(QUALITY_CHECK_CHUNK).collect();
        let mut evaluations: Vec<types::RolloutEvaluation> = Vec::new();
        let mut quality_check: Option<QualityCheck> = None;
        let total_chunks = chunks.len();
        for (part, chunk) in chunks.into_iter().enumerate() {
            let chunk_label = match (&opts.label, total_chunks) {
                (Some(l), n) if n > 1 => Some(format!("{l} {}/{n}", part + 1)),
                (Some(l), _) => Some(l.clone()),
                (None, _) => None,
            };
            let mut check = self
                .rollouts
                .quality_check(chunk, chunk_label.as_deref())
                .await?;
            check
                .wait(WaitOptions {
                    verbose: opts.verbose,
                    ..Default::default()
                })
                .await?;
            evaluations.extend(check.rollouts().await?);
            quality_check = Some(check);
        }
        let quality_check = quality_check.ok_or_else(|| {
            Error::Api("no rollouts were produced, so nothing could be judged".to_string())
        })?;
        let passed = evaluations.iter().filter(|e| e.passed).count();
        let pass_rate = if evaluations.is_empty() {
            None
        } else {
            Some(passed as f64 / evaluations.len() as f64)
        };
        Ok(EvalOutcome {
            prepare,
            dataset: target,
            results,
            quality_check,
            evaluations,
            pass_rate,
        })
    }

    /// The seeds that grounded a probe session, or `None` if unreadable.
    /// Grounding is baked into the representation once, at probe convergence, so
    /// a session's seed set is the exact cache key for everything downstream.
    async fn probe_seed_ids(&self, probe_session_id: &str) -> Result<Option<Vec<String>>> {
        match self
            .http
            .get::<types::ProbeSession>(&format!("/v1/probe-sessions/{probe_session_id}"), &[])
            .await
        {
            Ok(ps) => Ok(Some(ps.seed_ids)),
            Err(e) if e.status() == Some(404) => Ok(None),
            Err(e) => Err(e),
        }
    }

    /// Whether a probed artifact may be reused. `None` means the caller didn't
    /// ask about grounding, so reuse is unchanged — and short-circuits before any
    /// HTTP, which keeps existing callers' request sequences identical.
    async fn grounding_matches(
        &self,
        probe_session_id: &str,
        seed_ids: Option<&[String]>,
    ) -> Result<bool> {
        let Some(requested) = seed_ids else {
            return Ok(true);
        };
        match self.probe_seed_ids(probe_session_id).await? {
            Some(attached) => Ok(same_id_set(&attached, requested)),
            None => Ok(false),
        }
    }

    async fn probe_and_generate(
        &mut self,
        agent: &dyn ProbeAgent,
        opts: &PrepareOptions,
        mut reason: String,
        success_rate: Option<f64>,
        quality_check_id: Option<String>,
        force_probe: bool,
    ) -> Result<PrepareResult> {
        // Without a drift signal the agent hasn't been shown to change, so a
        // user model this session already probed is still good — skip the probe.
        let mut user_model: Option<types::UserModel> = None;
        if self.session_id.is_some() && quality_check_id.is_none() && !force_probe {
            let session_models = self.user_models.list(self.session_id.as_deref()).await?;
            if let Some(candidate) = session_models.last().cloned() {
                // ...but a model probed under different grounding would generate
                // from the wrong facts, and grounding can't be patched in after
                // convergence — only a fresh probe carries it.
                if self
                    .grounding_matches(&candidate.probe_session_id, opts.seed_ids.as_deref())
                    .await?
                {
                    user_model = Some(candidate);
                    reason += "; reusing session user model (no drift signal)";
                } else {
                    reason += "; requested grounding differs from the session user model's \
                               — re-probing";
                }
            }
        }
        let user_model = match user_model {
            Some(m) => m,
            None => {
                self.user_models
                    .create_from_probe(
                        agent,
                        ProbeOptions {
                            max_turns: opts.max_turns,
                            verbose: opts.verbose,
                            seed_ids: opts.seed_ids.clone(),
                        },
                    )
                    .await?
            }
        };
        let job = self
            .datasets
            .generate(
                &user_model.id,
                GenerateOptions {
                    count: Some(opts.count),
                    quality_check_id: quality_check_id.clone(),
                    guidance: None,
                },
            )
            .await?;
        let dataset = job
            .wait(WaitOptions {
                verbose: opts.verbose,
                ..Default::default()
            })
            .await?;
        Ok(PrepareResult {
            dataset,
            user_model,
            action: "generated".to_string(),
            reason,
            success_rate,
            quality_check_id,
        })
    }
}