a3s_code_core/llm/

structured.rs

//! Structured object generation from LLM output.
//!
//! Provides reliable JSON object generation with schema validation, automatic
//! repair, and streaming partial object support. Works across all providers by
//! selecting the best available mode (strict JSON schema, json_mode, tool-call,
//! or prompt-only).

use super::{LlmClient, Message, StreamEvent, TokenUsage, ToolDefinition};
use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio_util::sync::CancellationToken;

mod partial_json;
use partial_json::parse_partial_json;
#[cfg(test)]
use partial_json::try_parse_partial_json;

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// Mode selection for structured output generation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StructuredMode {
    /// Auto-select best mode based on provider capabilities.
    Auto,
    /// OpenAI native strict JSON schema (response_format.type = json_schema).
    Strict,
    /// OpenAI json_object mode (guarantees valid JSON, not schema-conformant).
    Json,
    /// Use tool-calling: inject a synthetic tool whose parameters IS the schema.
    /// Works on all providers that support tool use (Anthropic, OpenAI, etc).
    Tool,
    /// Prompt-only: append schema instructions to the prompt. Least reliable.
    Prompt,
}

/// Request specification for structured object generation.
#[derive(Debug, Clone)]
pub struct StructuredRequest {
    pub prompt: String,
    pub system: Option<String>,
    pub schema: Value,
    pub schema_name: String,
    pub schema_description: Option<String>,
    pub mode: StructuredMode,
    pub max_repair_attempts: u8,
}

/// Result of a successful structured generation.
#[derive(Debug, Clone, Serialize)]
pub struct StructuredResult {
    pub object: Value,
    pub raw_text: Option<String>,
    pub usage: TokenUsage,
    pub repair_rounds: u8,
    pub mode_used: StructuredMode,
}

/// Provider-native structured-output capability.
///
/// Each [`LlmClient`] reports this so the structured engine can request the
/// strongest enforcement the provider actually supports. Defaults to
/// [`NativeStructuredSupport::None`] for clients that don't override it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NativeStructuredSupport {
    /// No native enforcement — rely on prompt instructions + lenient extraction.
    None,
    /// Can force a specific tool call (Anthropic `tool_choice`, OpenAI function
    /// `tool_choice`). Guarantees the model emits the structured tool call
    /// instead of free-form prose.
    ForcedTool,
    /// Supports OpenAI-style `response_format` (`json_object` and
    /// `json_schema` + `strict`) in addition to forced tool calls.
    JsonSchema,
    /// Supports OpenAI-style `response_format: { type: "json_object" }`,
    /// but cannot combine it with a forced `tool_choice` (for example,
    /// DeepSeek reasoning models).  This is deliberately separate from
    /// [`JsonSchema`](Self::JsonSchema): callers can still get provider-side
    /// syntactic JSON guarantees without sending a request the model rejects.
    JsonObject,
}

/// A native `response_format` request for OpenAI-compatible providers.
#[derive(Debug, Clone, PartialEq)]
pub enum ResponseFormat {
    /// `{"type":"json_object"}` — guarantees syntactically valid JSON, but not
    /// schema conformance.
    JsonObject,
    /// `{"type":"json_schema","json_schema":{name,schema,strict:true}}` —
    /// parser-enforced schema conformance.
    JsonSchema { name: String, schema: Value },
}

/// Instruction telling a provider how to enforce structured output for a call.
///
/// Carries the union of intents; each provider honors what it supports and
/// ignores the rest (e.g. Anthropic has no `response_format`, so it only acts
/// on `force_tool`). [`Self::validation_schema`] is host-only metadata and must
/// never be serialized into a provider request. The default reproduces an
/// ordinary completion, which is why the trait's default `complete_structured`
/// implementation is behavior-preserving.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct StructuredDirective {
    /// Force the model to call exactly this tool (provider `tool_choice`).
    pub force_tool: Option<String>,
    /// Request a provider-native `response_format` (OpenAI-compatible only).
    pub response_format: Option<ResponseFormat>,
    /// Provider-facing response schema retained for composite-client stream
    /// validation, including prompt fallback and JSON-object modes where the
    /// provider directive itself does not carry a schema.
    pub validation_schema: Option<Value>,
}

/// Callback for streaming partial object snapshots.
pub type PartialObjectCallback = Box<dyn Fn(&Value) + Send>;

/// Provider-facing schema envelope.
///
/// Function/tool parameters are most reliable when the top-level schema is an
/// object. Inspired by Vercel AI SDK's `Output.array` / `Output.choice`
/// wrappers, A3S sends top-level arrays and scalar schemas inside a small object
/// envelope, then unwraps the validated value before returning it to callers.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SchemaEnvelope {
    Direct,
    Elements,
    Value,
}

impl SchemaEnvelope {
    fn for_schema(schema: &Value) -> Self {
        match schema_root_kind(schema, schema, &mut Vec::new(), 0) {
            Some(SchemaRootKind::Object) => Self::Direct,
            Some(SchemaRootKind::Array) => Self::Elements,
            Some(SchemaRootKind::Other) | None => Self::Value,
        }
    }

    fn response_schema(self, schema: &Value) -> Value {
        match self {
            Self::Direct => schema.clone(),
            Self::Elements => wrap_response_schema("elements", schema),
            Self::Value => wrap_response_schema("value", schema),
        }
    }

    fn unwrap_final(self, value: &Value) -> Option<Value> {
        match self {
            Self::Direct => Some(value.clone()),
            Self::Elements => value.get("elements").cloned(),
            Self::Value => value.get("value").cloned(),
        }
    }

    fn project_partial(self, value: &Value, repaired: bool) -> Option<Value> {
        match self {
            Self::Direct => Some(value.clone()),
            Self::Elements => {
                let mut elements = value.get("elements")?.as_array()?.clone();
                // A repaired parse may include a synthetic last element that was
                // closed only so the partial JSON can parse. Match Vercel's
                // array streaming behavior: publish only completed elements.
                if repaired && !elements.is_empty() {
                    elements.pop();
                }
                Some(Value::Array(elements))
            }
            Self::Value => value.get("value").cloned(),
        }
    }

    fn instruction(self) -> &'static str {
        match self {
            Self::Direct => "",
            Self::Elements => {
                "The provider-facing response schema wraps the requested array in an `elements` field. Follow that schema exactly; callers receive the unwrapped array."
            }
            Self::Value => {
                "The provider-facing response schema wraps the requested scalar/enum value in a `value` field. Follow that schema exactly; callers receive the unwrapped value."
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SchemaRootKind {
    Object,
    Array,
    Other,
}

fn schema_root_kind(
    schema: &Value,
    root: &Value,
    active_refs: &mut Vec<String>,
    depth: usize,
) -> Option<SchemaRootKind> {
    if depth > 64 {
        return None;
    }
    let object = schema.as_object()?;

    if let Some(kind) = object.get("type").and_then(schema_type_kind) {
        return Some(kind);
    }
    if let Some(value) = object.get("const") {
        return Some(value_kind(value));
    }
    if let Some(values) = object.get("enum").and_then(Value::as_array) {
        if let Some(kind) = common_value_kind(values) {
            return Some(kind);
        }
    }
    if let Some(reference) = object.get("$ref").and_then(Value::as_str) {
        if let Some(pointer) = reference.strip_prefix('#') {
            if !active_refs.iter().any(|active| active == reference) {
                if let Some(target) = root.pointer(pointer) {
                    active_refs.push(reference.to_string());
                    let kind = schema_root_kind(target, root, active_refs, depth + 1);
                    active_refs.pop();
                    if kind.is_some() {
                        return kind;
                    }
                }
            }
        }
    }
    if let Some(all_of) = object.get("allOf").and_then(Value::as_array) {
        if let Some(kind) = all_of
            .iter()
            .find_map(|branch| schema_root_kind(branch, root, active_refs, depth + 1))
        {
            return Some(kind);
        }
    }
    for keyword in ["anyOf", "oneOf"] {
        if let Some(branches) = object.get(keyword).and_then(Value::as_array) {
            let kinds = branches
                .iter()
                .map(|branch| schema_root_kind(branch, root, active_refs, depth + 1))
                .collect::<Option<Vec<_>>>();
            if let Some(kinds) = kinds {
                if let Some(first) = kinds.first().copied() {
                    if kinds.iter().all(|kind| *kind == first) {
                        return Some(first);
                    }
                }
            }
        }
    }

    // Preserve the established object-schema behavior for schemas that rely
    // on object-only keywords without an explicit `type`.
    if ["properties", "required", "additionalProperties"]
        .iter()
        .any(|keyword| object.contains_key(*keyword))
    {
        return Some(SchemaRootKind::Object);
    }
    None
}

fn schema_type_kind(value: &Value) -> Option<SchemaRootKind> {
    match value {
        Value::String(value) => Some(type_name_kind(value)),
        Value::Array(values) if values.len() == 1 => values[0].as_str().map(type_name_kind),
        _ => None,
    }
}

fn type_name_kind(value: &str) -> SchemaRootKind {
    match value {
        "object" => SchemaRootKind::Object,
        "array" => SchemaRootKind::Array,
        _ => SchemaRootKind::Other,
    }
}

fn value_kind(value: &Value) -> SchemaRootKind {
    match value {
        Value::Object(_) => SchemaRootKind::Object,
        Value::Array(_) => SchemaRootKind::Array,
        _ => SchemaRootKind::Other,
    }
}

fn common_value_kind(values: &[Value]) -> Option<SchemaRootKind> {
    let first = values.first().map(value_kind)?;
    values
        .iter()
        .all(|value| value_kind(value) == first)
        .then_some(first)
}

fn wrap_response_schema(field: &str, schema: &Value) -> Value {
    let mut embedded = schema.clone();
    let mut wrapper = serde_json::json!({
        "type": "object",
        "required": [field],
        "additionalProperties": false,
        "properties": {}
    });

    // A local reference such as `#/$defs/item` resolves from the provider-
    // facing document root. Hoist root definitions when the requested schema
    // must be wrapped so those references retain their original meaning.
    if let Some(embedded_object) = embedded.as_object_mut() {
        for keyword in ["$defs", "definitions"] {
            if let Some(definitions) = embedded_object.remove(keyword) {
                wrapper[keyword] = definitions;
            }
        }
    }
    wrapper["properties"][field] = embedded;
    wrapper
}

// ---------------------------------------------------------------------------
// Core generation: blocking (non-streaming)
// ---------------------------------------------------------------------------

/// Generate a structured JSON object using the given LLM client.
///
/// Selects the best mode based on `req.mode`, calls the LLM, validates against
/// the schema, and retries with repair prompts if validation fails.
pub async fn generate_blocking(
    client: &dyn LlmClient,
    req: &StructuredRequest,
) -> Result<StructuredResult> {
    generate_blocking_with_cancellation(client, req, CancellationToken::new()).await
}

/// Generate a structured object while honoring the caller's cancellation token.
///
/// This is the explicit model-call boundary for non-streaming structured
/// generation. The client remains the only provider gateway; cancellation is
/// enforced around every initial and repair call without mutating the caller's
/// token.
pub async fn generate_blocking_with_cancellation(
    client: &dyn LlmClient,
    req: &StructuredRequest,
    cancellation: CancellationToken,
) -> Result<StructuredResult> {
    let mode = resolve_mode(req.mode, client.native_structured_support());
    let envelope = SchemaEnvelope::for_schema(&req.schema);
    let mut messages = build_initial_messages(req, mode);
    let system = build_system_prompt(req, mode);
    let tools = build_tools(req, mode);
    let directive = build_directive(req, mode);

    let mut total_usage = TokenUsage::default();
    let mut repair_rounds: u8 = 0;

    loop {
        let resp = tokio::select! {
            biased;
            _ = cancellation.cancelled() => bail!("Operation cancelled by user"),
            response = client.complete_structured(&messages, Some(&system), &tools, &directive) => response,
        }
        .context("LLM call failed during structured generation")?;

        accumulate_usage(&mut total_usage, &resp.usage);

        // Mine the object from every place a model might have parked it (tool call,
        // text content, AND the reasoning channel), trying each balanced JSON
        // candidate against the schema. Reasoning models routinely leave `content`
        // empty and emit the object inside `reasoning`, so without the reasoning
        // fallback generate_object failed with "no structured output" across models.
        let candidates = extract_raw_candidates(&resp.message, mode);
        let resolution = resolve_structured(&candidates, &req.schema, envelope);

        if let Some((value, raw)) = resolution.valid {
            return Ok(StructuredResult {
                object: value,
                raw_text: Some(raw),
                usage: total_usage,
                repair_rounds,
                mode_used: mode,
            });
        }

        if repair_rounds >= req.max_repair_attempts {
            return Err(match resolution.invalid {
                Some((_, errors)) => anyhow::anyhow!(
                    "Structured output failed schema validation after {} repair attempts. Errors: {}",
                    repair_rounds,
                    errors.join("; ")
                ),
                None => anyhow::anyhow!(
                    "Structured output parsing failed after {} repair attempts: no JSON object found in tool call, text content, or reasoning channel",
                    repair_rounds
                ),
            });
        }

        repair_rounds += 1;
        let (repair_msg, raw_for_ctx) = match resolution.invalid {
            Some((raw, errors)) => (build_repair_message(&raw, &errors), raw),
            None => {
                let raw = resolution.raw_seen.unwrap_or_default();
                (build_parse_failure_repair(&raw), raw)
            }
        };
        append_repair_context(
            &mut messages,
            &resp.message,
            &repair_msg,
            mode,
            &raw_for_ctx,
        );
    }
}

// ---------------------------------------------------------------------------
// Core generation: streaming
// ---------------------------------------------------------------------------

/// Generate a structured JSON object with streaming partial updates.
///
/// Calls `on_partial` with progressively more complete partial objects as tokens
/// arrive. Returns the final validated object.
///
/// A streamed first attempt may be followed by bounded non-streaming repair
/// calls when `max_repair_attempts` is non-zero. Repair calls publish only the
/// final corrected object, avoiding a second misleading partial stream.
pub async fn generate_streaming(
    client: &dyn LlmClient,
    req: &StructuredRequest,
    on_partial: PartialObjectCallback,
) -> Result<StructuredResult> {
    generate_streaming_with_cancellation(client, req, on_partial, CancellationToken::new()).await
}

/// Generate a structured object with streaming partial updates while honoring
/// the caller's cancellation token.
///
/// A child token is passed to the provider so cancelling this operation can
/// abort provider I/O without ever cancelling a host-owned token. Repair calls
/// use the same cancellation boundary as the initial stream.
pub async fn generate_streaming_with_cancellation(
    client: &dyn LlmClient,
    req: &StructuredRequest,
    on_partial: PartialObjectCallback,
    cancellation: CancellationToken,
) -> Result<StructuredResult> {
    let mode = resolve_mode(req.mode, client.native_structured_support());
    let envelope = SchemaEnvelope::for_schema(&req.schema);
    let mut messages = build_initial_messages(req, mode);
    let system = build_system_prompt(req, mode);
    let tools = build_tools(req, mode);
    let directive = build_directive(req, mode);

    let provider_cancellation = cancellation.child_token();
    let mut rx = tokio::select! {
        biased;
        _ = cancellation.cancelled() => bail!("Operation cancelled by user"),
        response = client.complete_streaming_structured(
            &messages,
            Some(&system),
            &tools,
            &directive,
            provider_cancellation.clone(),
        ) => response,
    }
    .context("LLM streaming call failed during structured generation")?;

    let mut json_buffer = String::new();
    let mut last_valid_partial: Option<Value> = None;
    let mut final_response: Option<super::LlmResponse> = None;
    let mut last_parse_len: usize = 0;
    let mut complete_candidate: Option<(Value, String, tokio::time::Instant)> = None;
    // Minimum bytes of new data before attempting a partial parse (reduces CPU)
    const PARSE_THRESHOLD: usize = 8;
    // Well-behaved providers send Done immediately after the complete object.
    // A short grace preserves their final usage metadata while preventing an
    // otherwise valid result from hanging on a compatible endpoint that never
    // terminates its stream.
    const DONE_GRACE: std::time::Duration = std::time::Duration::from_millis(250);
    loop {
        let event = if let Some((_, _, deadline)) = complete_candidate.as_ref() {
            tokio::select! {
                biased;
                _ = cancellation.cancelled() => bail!("Operation cancelled by user"),
                event = rx.recv() => event,
                _ = tokio::time::sleep_until(*deadline) => {
                    let Some(candidate) = complete_candidate.take() else {
                        continue;
                    };
                    let (value, raw_text, _) = candidate;
                    provider_cancellation.cancel();
                    on_partial(&value);
                    return Ok(StructuredResult {
                        object: value,
                        raw_text: Some(raw_text),
                        usage: TokenUsage::default(),
                        repair_rounds: 0,
                        mode_used: mode,
                    });
                }
            }
        } else {
            tokio::select! {
                biased;
                _ = cancellation.cancelled() => bail!("Operation cancelled by user"),
                event = rx.recv() => event,
            }
        };
        let Some(event) = event else {
            if let Some((value, raw_text, _)) = complete_candidate.take() {
                provider_cancellation.cancel();
                on_partial(&value);
                return Ok(StructuredResult {
                    object: value,
                    raw_text: Some(raw_text),
                    usage: TokenUsage::default(),
                    repair_rounds: 0,
                    mode_used: mode,
                });
            }
            break;
        };
        match event {
            StreamEvent::ToolUseInputDelta { delta, .. } if mode == StructuredMode::Tool => {
                if final_response.is_some() {
                    continue;
                }
                json_buffer.push_str(&delta);
                if json_buffer.len() - last_parse_len >= PARSE_THRESHOLD {
                    if let Some(partial) = parse_partial_json(&json_buffer) {
                        if let Some(projected) =
                            envelope.project_partial(&partial.value, partial.repaired)
                        {
                            if last_valid_partial.as_ref() != Some(&projected) {
                                on_partial(&projected);
                                last_valid_partial = Some(projected);
                            }
                        }
                    }
                    last_parse_len = json_buffer.len();
                }
                if complete_candidate.is_none() && (delta.contains('}') || delta.contains(']')) {
                    complete_candidate = resolve_structured(
                        std::slice::from_ref(&json_buffer),
                        &req.schema,
                        envelope,
                    )
                    .valid
                    .map(|(value, raw_text)| {
                        (value, raw_text, tokio::time::Instant::now() + DONE_GRACE)
                    });
                }
            }
            StreamEvent::TextDelta(delta) if mode != StructuredMode::Tool => {
                if final_response.is_some() {
                    continue;
                }
                json_buffer.push_str(&delta);
                if json_buffer.len() - last_parse_len >= PARSE_THRESHOLD {
                    if let Some(json_start) = find_json_start(&json_buffer) {
                        let candidate = &json_buffer[json_start..];
                        if let Some(partial) = parse_partial_json(candidate) {
                            if let Some(projected) =
                                envelope.project_partial(&partial.value, partial.repaired)
                            {
                                if last_valid_partial.as_ref() != Some(&projected) {
                                    on_partial(&projected);
                                    last_valid_partial = Some(projected);
                                }
                            }
                        }
                    }
                    last_parse_len = json_buffer.len();
                }
                if complete_candidate.is_none() && (delta.contains('}') || delta.contains(']')) {
                    complete_candidate = resolve_structured(
                        std::slice::from_ref(&json_buffer),
                        &req.schema,
                        envelope,
                    )
                    .valid
                    .map(|(value, raw_text)| {
                        (value, raw_text, tokio::time::Instant::now() + DONE_GRACE)
                    });
                }
            }
            StreamEvent::Done(resp) => {
                final_response = Some(resp);
                break;
            }
            _ => {}
        }
    }

    provider_cancellation.cancel();
    let mut resp = final_response.context("Stream ended without Done event")?;
    let mut total_usage = TokenUsage::default();
    accumulate_usage(&mut total_usage, &resp.usage);
    let mut repair_rounds = 0u8;
    // Same multi-source resolution as the blocking path: the final message may carry
    // the object in the tool call, the text content, or the reasoning channel.
    let mut resolution = resolve_structured(
        &extract_raw_candidates(&resp.message, mode),
        &req.schema,
        envelope,
    );
    let (value, raw_text) = loop {
        if let Some(valid) = resolution.valid.take() {
            break valid;
        }

        if repair_rounds >= req.max_repair_attempts {
            return Err(match resolution.invalid {
                Some((_, errors)) => anyhow::anyhow!(
                    "Streamed structured output failed schema validation after {} repair attempts: {}",
                    repair_rounds,
                    errors.join("; ")
                ),
                None => anyhow::anyhow!(
                    "Streamed output produced no parseable JSON object after {} repair attempts (checked tool call, text content, and reasoning channel)",
                    repair_rounds
                ),
            });
        }

        repair_rounds += 1;
        let (repair_message, raw_for_context) = match resolution.invalid.take() {
            Some((raw, errors)) => (build_repair_message(&raw, &errors), raw),
            None => {
                let raw = resolution.raw_seen.take().unwrap_or_default();
                (build_parse_failure_repair(&raw), raw)
            }
        };
        append_repair_context(
            &mut messages,
            &resp.message,
            &repair_message,
            mode,
            &raw_for_context,
        );
        resp = tokio::select! {
            biased;
            _ = cancellation.cancelled() => bail!("Operation cancelled by user"),
            response = client.complete_structured(&messages, Some(&system), &tools, &directive) => response,
        }
        .context("LLM call failed while repairing streamed structured output")?;
        accumulate_usage(&mut total_usage, &resp.usage);
        resolution = resolve_structured(
            &extract_raw_candidates(&resp.message, mode),
            &req.schema,
            envelope,
        );
    };

    // Emit final complete object
    on_partial(&value);

    Ok(StructuredResult {
        object: value,
        raw_text: Some(raw_text),
        usage: total_usage,
        repair_rounds,
        mode_used: mode,
    })
}

// ---------------------------------------------------------------------------
// JSON extraction and parsing
// ---------------------------------------------------------------------------

/// Extract a JSON value from potentially dirty LLM output.
///
/// Handles: raw JSON, markdown code fences, leading/trailing prose.
pub fn extract_json_value(text: &str) -> Result<Value> {
    let trimmed = text.trim();

    // 1. Direct parse
    if let Ok(v) = serde_json::from_str::<Value>(trimmed) {
        if v.is_object() || v.is_array() {
            return Ok(v);
        }
    }

    // 2. Strip markdown code fence
    if let Some(inner) = strip_code_fence(trimmed) {
        if let Ok(v) = serde_json::from_str::<Value>(inner.trim()) {
            if v.is_object() || v.is_array() {
                return Ok(v);
            }
        }
    }

    // 3. Find balanced JSON substring (first { to matching })
    if let Some(candidate) = find_balanced_json_object(trimmed) {
        if let Ok(v) = serde_json::from_str::<Value>(candidate) {
            return Ok(v);
        }
    }

    // 4. Try array
    if let Some(candidate) = find_balanced_json_array(trimmed) {
        if let Ok(v) = serde_json::from_str::<Value>(candidate) {
            return Ok(v);
        }
    }

    bail!("No valid JSON object found in LLM output")
}

/// Strip ```json ... ``` or ``` ... ``` fences.
fn strip_code_fence(text: &str) -> Option<&str> {
    let start_patterns = ["```json\n", "```json\r\n", "```\n", "```\r\n"];
    for pat in &start_patterns {
        if let Some(rest) = text.strip_prefix(pat) {
            // Find closing fence
            if let Some(end) = rest.rfind("```") {
                return Some(&rest[..end]);
            }
        }
    }
    // Also handle inline: ```json{...}```
    if let Some(inner) = text.strip_prefix("```json") {
        if let Some(end) = inner.rfind("```") {
            return Some(inner[..end].trim());
        }
    }
    if let Some(inner) = text.strip_prefix("```") {
        if let Some(end) = inner.rfind("```") {
            return Some(inner[..end].trim());
        }
    }
    None
}

/// Find the first balanced `{...}` substring using bracket counting.
fn find_balanced_json_object(text: &str) -> Option<&str> {
    find_balanced(text, '{', '}')
}

/// Find the first balanced `[...]` substring.
fn find_balanced_json_array(text: &str) -> Option<&str> {
    find_balanced(text, '[', ']')
}

fn find_balanced(text: &str, open: char, close: char) -> Option<&str> {
    find_balanced_range(text, open, close).map(|(start, end)| &text[start..end])
}

/// Byte range `[start, end)` of the first balanced `open..close` substring (quote-aware).
fn find_balanced_range(text: &str, open: char, close: char) -> Option<(usize, usize)> {
    let bytes = text.as_bytes();
    let open_byte = open as u8;
    let close_byte = close as u8;

    // Find the first unquoted occurrence of `open`
    let mut in_string = false;
    let mut escape_next = false;
    let mut start = None;

    for (i, &b) in bytes.iter().enumerate() {
        if escape_next {
            escape_next = false;
            continue;
        }
        match b {
            b'\\' if in_string => escape_next = true,
            b'"' => in_string = !in_string,
            _ if in_string => {}
            _ if b == open_byte => {
                start = Some(i);
                break;
            }
            _ => {}
        }
    }

    let start = start?;
    let mut depth = 0i32;
    in_string = false;
    escape_next = false;

    for (i, &b) in bytes[start..].iter().enumerate() {
        if escape_next {
            escape_next = false;
            continue;
        }
        match b {
            b'\\' if in_string => escape_next = true,
            b'"' => in_string = !in_string,
            _ if in_string => {}
            _ if b == open_byte => depth += 1,
            _ if b == close_byte => {
                depth -= 1;
                if depth == 0 {
                    return Some((start, start + i + 1));
                }
            }
            _ => {}
        }
    }
    None
}

/// Every top-level balanced `open..close` substring, in document order.
///
/// Reasoning traces often contain several objects (worked examples, partial drafts)
/// before the final answer, so callers validate each against the schema and keep the
/// one that fits rather than blindly trusting the first `{...}`.
fn find_all_balanced(text: &str, open: char, close: char) -> Vec<String> {
    let mut out = Vec::new();
    let mut base = 0usize;
    while base < text.len() {
        match find_balanced_range(&text[base..], open, close) {
            Some((start, end)) => {
                out.push(text[base + start..base + end].to_string());
                base += end;
            }
            None => break,
        }
    }
    out
}

/// Find the byte offset where JSON content starts in a text stream.
/// Skips leading prose/whitespace to find `{` or `[` that isn't inside a string.
fn find_json_start(text: &str) -> Option<usize> {
    // Skip past code fence markers if present
    let (search_text, offset) = if let Some(rest) = text.strip_prefix("```json") {
        (rest, 7)
    } else if let Some(rest) = text.strip_prefix("```") {
        (rest, 3)
    } else {
        (text, 0)
    };

    let mut in_string = false;
    let mut escape_next = false;
    for (i, &b) in search_text.as_bytes().iter().enumerate() {
        if escape_next {
            escape_next = false;
            continue;
        }
        match b {
            b'\\' if in_string => {
                escape_next = true;
            }
            b'"' => {
                in_string = !in_string;
            }
            b'{' | b'[' if !in_string => {
                return Some(offset + i);
            }
            _ => {}
        }
    }
    None
}

// ---------------------------------------------------------------------------
// Schema validation
// ---------------------------------------------------------------------------

/// Validate a JSON value against a JSON Schema.
/// Returns Ok(()) on success, or a list of human-readable error strings.
fn validate_against_schema(value: &Value, schema: &Value) -> Result<(), Vec<String>> {
    // Structured-output schemas are host/model input, so compilation is kept
    // entirely in-memory: the dependency is built without HTTP/file resolvers.
    // Local `$ref` / `$defs`, composition keywords, conditional schemas, and
    // exact `oneOf` semantics are handled by the standards-compliant validator.
    let validator = jsonschema::draft202012::options()
        .build(schema)
        .map_err(|error| vec![format!("invalid JSON Schema: {error}")])?;
    let errors = validator
        .iter_errors(value)
        .map(|error| {
            let path = error.instance_path().to_string();
            if path.is_empty() {
                format!("$: {error}")
            } else {
                format!("{path}: {error}")
            }
        })
        .collect::<Vec<_>>();
    if errors.is_empty() {
        Ok(())
    } else {
        Err(errors)
    }
}

/// Return whether a streamed provider payload is already one complete JSON
/// value that conforms to the provider-facing response schema.
///
/// Composite clients can use this to keep candidate streams isolated until a
/// response is safe to replay, while still accepting a schema-complete object
/// from endpoints that omit their terminal stream event. The strict direct
/// parse deliberately rejects prose, code fences, and trailing data; the main
/// structured engine remains responsible for its broader post-response repair
/// behavior.
pub fn is_complete_streamed_value(raw: &str, response_schema: &Value) -> bool {
    serde_json::from_str::<Value>(raw)
        .ok()
        .is_some_and(|value| validate_against_schema(&value, response_schema).is_ok())
}

// ---------------------------------------------------------------------------
// Message/prompt construction helpers
// ---------------------------------------------------------------------------

/// Resolve the requested mode against the provider's native capability.
///
/// Prefer native enforcement only when the client explicitly reports support.
/// Unknown OpenAI-compatible endpoints can hang when sent `tool_choice` or
/// `response_format`, so unsupported requests degrade to prompt+schema parsing
/// instead of optimistic native parameters.
fn resolve_mode(requested: StructuredMode, support: NativeStructuredSupport) -> StructuredMode {
    match (requested, support) {
        (StructuredMode::Prompt, _) => StructuredMode::Prompt,
        (StructuredMode::Strict, NativeStructuredSupport::JsonSchema) => StructuredMode::Strict,
        (StructuredMode::Json, NativeStructuredSupport::JsonSchema) => StructuredMode::Json,
        (StructuredMode::Auto | StructuredMode::Tool, NativeStructuredSupport::JsonSchema) => {
            StructuredMode::Tool
        }
        (
            StructuredMode::Auto
            | StructuredMode::Tool
            | StructuredMode::Strict
            | StructuredMode::Json,
            NativeStructuredSupport::ForcedTool,
        ) => StructuredMode::Tool,
        (
            StructuredMode::Auto
            | StructuredMode::Tool
            | StructuredMode::Strict
            | StructuredMode::Json,
            NativeStructuredSupport::JsonObject,
        ) => StructuredMode::Json,
        (
            StructuredMode::Auto
            | StructuredMode::Tool
            | StructuredMode::Strict
            | StructuredMode::Json,
            NativeStructuredSupport::None,
        ) => StructuredMode::Prompt,
    }
}

/// Build the provider directive for an already-resolved mode.
fn build_directive(req: &StructuredRequest, mode: StructuredMode) -> StructuredDirective {
    let response_schema = SchemaEnvelope::for_schema(&req.schema).response_schema(&req.schema);
    let mut directive = match mode {
        StructuredMode::Tool => StructuredDirective {
            force_tool: Some(format!("emit_{}", req.schema_name)),
            response_format: None,
            validation_schema: None,
        },
        StructuredMode::Strict => StructuredDirective {
            force_tool: None,
            response_format: Some(ResponseFormat::JsonSchema {
                name: req.schema_name.clone(),
                schema: response_schema.clone(),
            }),
            validation_schema: None,
        },
        StructuredMode::Json => StructuredDirective {
            force_tool: None,
            response_format: Some(ResponseFormat::JsonObject),
            validation_schema: None,
        },
        StructuredMode::Auto | StructuredMode::Prompt => StructuredDirective::default(),
    };
    directive.validation_schema = Some(response_schema);
    directive
}

fn build_initial_messages(req: &StructuredRequest, mode: StructuredMode) -> Vec<Message> {
    let envelope = SchemaEnvelope::for_schema(&req.schema);
    let response_schema = envelope.response_schema(&req.schema);
    let envelope_instruction = envelope.instruction();
    match mode {
        StructuredMode::Tool => {
            // For tool mode, the prompt is the user message; the LLM will respond
            // with a tool call whose input is the structured object.
            vec![Message::user(&req.prompt)]
        }
        StructuredMode::Prompt | StructuredMode::Json => {
            // Prompt mode and json_object mode both need the schema in the prompt:
            // json_object only guarantees *syntactic* validity, so the model still
            // has to be told the shape it should produce.
            let augmented = format!(
                "{}\n\n{}{}\n\nYou MUST respond with ONLY a valid JSON object (no markdown, no explanation) that conforms to this JSON Schema:\n\n```json\n{}\n```",
                req.prompt,
                envelope_instruction,
                if envelope_instruction.is_empty() { "" } else { "\n" },
                serde_json::to_string_pretty(&response_schema).unwrap_or_default()
            );
            vec![Message::user(&augmented)]
        }
        _ => {
            // Strict mode: the schema constraint is enforced by the provider via
            // response_format.json_schema, so the user message is just the prompt.
            vec![Message::user(&req.prompt)]
        }
    }
}

fn build_system_prompt(req: &StructuredRequest, mode: StructuredMode) -> String {
    let base = req.system.as_deref().unwrap_or("");
    let envelope_instruction = SchemaEnvelope::for_schema(&req.schema).instruction();

    match mode {
        StructuredMode::Tool => {
            format!(
                "{}{}You MUST respond by calling the `emit_{}` tool exactly once with a valid argument matching the schema. Do not output any text outside the tool call.{}{}",
                base,
                if base.is_empty() { "" } else { "\n\n" },
                req.schema_name,
                if envelope_instruction.is_empty() { "" } else { "\n\n" },
                envelope_instruction
            )
        }
        StructuredMode::Prompt | StructuredMode::Json => {
            format!(
                "{}{}You are a structured data extraction assistant. Always respond with valid JSON only, no markdown fences, no explanation text.{}{}",
                base,
                if base.is_empty() { "" } else { "\n\n" },
                if envelope_instruction.is_empty() { "" } else { "\n\n" },
                envelope_instruction,
            )
        }
        _ => base.to_string(),
    }
}

fn build_tools(req: &StructuredRequest, mode: StructuredMode) -> Vec<ToolDefinition> {
    match mode {
        StructuredMode::Tool => {
            vec![ToolDefinition {
                name: format!("emit_{}", req.schema_name),
                description: req
                    .schema_description
                    .clone()
                    .unwrap_or_else(|| format!("Emit a structured {} object", req.schema_name)),
                parameters: SchemaEnvelope::for_schema(&req.schema).response_schema(&req.schema),
            }]
        }
        _ => vec![],
    }
}

/// Outcome of mining a response for the structured object across all candidate sources.
struct StructuredResolution {
    /// A schema-valid object plus the raw source string it came from.
    valid: Option<(Value, String)>,
    /// First parseable-but-schema-invalid object source + its validation errors,
    /// used to build a targeted repair prompt.
    invalid: Option<(String, Vec<String>)>,
    /// First non-empty raw candidate, shown verbatim in a parse-failure repair prompt.
    raw_seen: Option<String>,
}

/// Append `s` to `out` if it is non-empty and not already present (trimmed, deduped).
fn push_candidate(out: &mut Vec<String>, s: String) {
    let trimmed = s.trim();
    if !trimmed.is_empty() && !out.iter().any(|c| c == trimmed) {
        out.push(trimmed.to_string());
    }
}

/// Ordered raw strings to mine for the structured object, most authoritative first:
/// tool-call arguments, then text content, then the reasoning channel.
///
/// The reasoning fallback is the crux of the cross-model fix: reasoning models
/// (GLM/zhipu, DeepSeek-R1, kimi…) frequently emit the final object inside
/// `reasoning` with `content` empty and no tool call. Earlier extraction only looked
/// at the tool call / text, so those models yielded an empty string and the whole
/// generate_object failed even though a perfectly good object was produced.
fn extract_raw_candidates(message: &super::Message, mode: StructuredMode) -> Vec<String> {
    let mut out: Vec<String> = Vec::new();
    if mode == StructuredMode::Tool {
        if let Some(call) = message.tool_calls().first() {
            push_candidate(
                &mut out,
                serde_json::to_string(&call.args).unwrap_or_default(),
            );
        }
    }
    push_candidate(&mut out, message.text());
    if let Some(reasoning) = message.reasoning_content.as_deref() {
        push_candidate(&mut out, reasoning.to_string());
    }
    out
}

/// Every JSON object/array value mineable from possibly-dirty text, in document order
/// (direct parse, code fences, then all balanced `{...}` / `[...]`). Deduped.
#[cfg(test)]
fn extract_all_json_values(text: &str) -> Vec<Value> {
    extract_json_candidates(text, false)
}

/// Every JSON value mineable from possibly-dirty text for schema-aware structured
/// resolution. When `include_direct_scalars` is true, direct raw/fenced scalar JSON
/// is retained so top-level scalar schemas can recover non-enveloped model output.
fn extract_json_candidates(text: &str, include_direct_scalars: bool) -> Vec<Value> {
    let trimmed = text.trim();
    let mut values: Vec<Value> = Vec::new();
    let consider = |candidate: &str, values: &mut Vec<Value>, allow_scalar: bool| {
        if let Ok(v) = serde_json::from_str::<Value>(candidate.trim()) {
            if (v.is_object() || v.is_array() || allow_scalar) && !values.contains(&v) {
                values.push(v);
            }
        }
    };
    consider(trimmed, &mut values, include_direct_scalars);
    if let Some(inner) = strip_code_fence(trimmed) {
        consider(inner, &mut values, include_direct_scalars);
    }
    for candidate in find_all_balanced(trimmed, '{', '}') {
        consider(&candidate, &mut values, false);
    }
    for candidate in find_all_balanced(trimmed, '[', ']') {
        consider(&candidate, &mut values, false);
    }
    values
}

/// Try every raw candidate × every JSON value it yields against the schema; return the
/// first schema-valid value, else the best parseable-but-invalid value (for repair).
fn resolve_structured(
    candidates: &[String],
    schema: &Value,
    envelope: SchemaEnvelope,
) -> StructuredResolution {
    let mut invalid: Option<(String, Vec<String>)> = None;
    let mut raw_seen: Option<String> = None;
    let response_schema = envelope.response_schema(schema);
    for raw in candidates {
        if raw_seen.is_none() && !raw.trim().is_empty() {
            raw_seen = Some(raw.clone());
        }
        for value in extract_json_candidates(raw, envelope == SchemaEnvelope::Value) {
            match validate_against_schema(&value, schema) {
                Ok(()) => {
                    return StructuredResolution {
                        valid: Some((value, raw.clone())),
                        invalid,
                        raw_seen,
                    };
                }
                Err(errors) => {
                    if invalid.is_none() {
                        invalid = Some((raw.clone(), errors));
                    }
                }
            }

            if envelope != SchemaEnvelope::Direct {
                match validate_against_schema(&value, &response_schema) {
                    Ok(()) => {
                        if let Some(unwrapped) = envelope.unwrap_final(&value) {
                            match validate_against_schema(&unwrapped, schema) {
                                Ok(()) => {
                                    return StructuredResolution {
                                        valid: Some((unwrapped, raw.clone())),
                                        invalid,
                                        raw_seen,
                                    };
                                }
                                Err(errors) => {
                                    if invalid.is_none() {
                                        invalid = Some((raw.clone(), errors));
                                    }
                                }
                            }
                        } else if invalid.is_none() {
                            invalid = Some((
                                raw.clone(),
                                vec!["$: response envelope was missing the expected value field"
                                    .to_string()],
                            ));
                        }
                    }
                    Err(errors) => {
                        if invalid.is_none() {
                            invalid = Some((raw.clone(), errors));
                        }
                    }
                }
            }
        }
    }
    StructuredResolution {
        valid: None,
        invalid,
        raw_seen,
    }
}

/// Extract the first JSON value from possibly dirty model text that validates
/// against `schema`.
///
/// This is the local fast path for callers that already asked an agent to
/// produce structured output. It accepts direct JSON, fenced JSON, and a
/// balanced JSON value embedded in prose, while preserving the same schema
/// and envelope semantics used by [`generate_blocking`]. Callers can fall back
/// to an LLM repair pass only when this returns `None`.
pub(crate) fn parse_validated_output(text: &str, schema: &Value) -> Option<Value> {
    resolve_structured(
        &[text.to_string()],
        schema,
        SchemaEnvelope::for_schema(schema),
    )
    .valid
    .map(|(value, _)| value)
}

/// UTF-8-safe truncation to at most `max` bytes (never splits a multibyte char —
/// repair prompts echo arbitrary model output, including CJK).
fn truncate_utf8(s: &str, max: usize) -> &str {
    if s.len() <= max {
        return s;
    }
    let mut end = max;
    while end > 0 && !s.is_char_boundary(end) {
        end -= 1;
    }
    &s[..end]
}

/// Repair prompt for when nothing parseable was produced at all.
fn build_parse_failure_repair(raw_text: &str) -> String {
    if raw_text.trim().is_empty() {
        return "Your previous response contained no JSON. Respond with ONLY a single valid JSON object that matches the schema — no prose, no markdown, no analysis, and put the object in your reply content (not in a thinking/reasoning aside).".to_string();
    }
    format!(
        "Your previous output could not be parsed as a JSON object:\n\n{}\n\nReturn ONLY a single valid JSON object matching the schema — no prose, no markdown.",
        truncate_utf8(raw_text, 2000)
    )
}

fn build_repair_message(raw_text: &str, errors: &[String]) -> String {
    // Truncate raw output in repair message to avoid blowing context
    let truncated_raw = if raw_text.len() > 2000 {
        format!(
            "{}...[truncated, {} bytes total]",
            truncate_utf8(raw_text, 2000),
            raw_text.len()
        )
    } else {
        raw_text.to_string()
    };
    format!(
        "Your previous output failed schema validation:\n\n{}\n\nValidation errors:\n{}\n\nPlease return ONLY a corrected JSON object that fixes these errors. No explanation, no markdown.",
        truncated_raw,
        errors.iter().map(|e| format!("- {}", e)).collect::<Vec<_>>().join("\n")
    )
}

fn accumulate_usage(total: &mut TokenUsage, delta: &TokenUsage) {
    total.prompt_tokens += delta.prompt_tokens;
    total.completion_tokens += delta.completion_tokens;
    total.total_tokens += delta.total_tokens;
}

/// Append repair context to the message history, respecting conversation structure.
///
/// In tool mode, the LLM returned a tool_use block. The correct follow-up is:
///   assistant (tool_use) → user (tool_result with error) → assistant (retry)
/// In text modes, it's simply:
///   assistant (text) → user (repair request) → assistant (retry)
fn append_repair_context(
    messages: &mut Vec<Message>,
    assistant_msg: &Message,
    repair_text: &str,
    mode: StructuredMode,
    _raw_text: &str,
) {
    if mode == StructuredMode::Tool {
        // Push the original assistant message (with tool_use block intact)
        messages.push(assistant_msg.clone());
        // Find the tool_use ID to construct a proper tool_result
        let tool_use_id = assistant_msg
            .tool_calls()
            .first()
            .map(|tc| tc.id.clone())
            .unwrap_or_else(|| "unknown".to_string());
        // Return the error as a tool_result so the conversation stays valid
        messages.push(Message::tool_result(&tool_use_id, repair_text, true));
    } else {
        // Text modes: push assistant text then user repair request
        messages.push(assistant_msg.clone());
        messages.push(Message::user(repair_text));
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
#[path = "structured_tests.rs"]
mod structured_tests;