audacity-sdk 0.2.0

//! § 3 — Request/response mapping between the canonical Bedrock-shaped API
//! and the OpenAI wire protocol spoken by the Audacity gateway.

use serde::Deserialize;
use serde_json::{json, Map, Value};

use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use base64::Engine as _;

use crate::error::{classify_by_code, extract_error_fields, ErrorDetails};
use crate::types::{
    ContentBlock, ContentBlockDeltaEvent, ContentBlockDeltaPayload, ContentBlockStartEvent,
    ContentBlockStartPayload, ContentBlockStartToolUse, ContentBlockStopEvent, ConversationRole,
    ConverseOutput, ConverseOutputEnum, ConverseStreamOutput, ImageBlock, ImageSource,
    InferenceConfiguration, Message, MessageStartEvent, MessageStopEvent, MetadataEvent, Metrics,
    StopReason, TokenUsage, ToolChoice, ToolConfiguration, ToolResultContent,
};
use crate::{ConverseInput, Error};

// ── Input → OpenAI request body ───────────────────────────────────────────────

pub(crate) fn build_request_body(input: &ConverseInput, stream: bool) -> Value {
    let mut body = Map::new();

    // 1. modelId → model
    body.insert("model".into(), json!(input.model_id));

    // 2 & 3. messages
    let messages = build_messages(input);
    body.insert("messages".into(), Value::Array(messages));

    // 4. inferenceConfig
    if let Some(cfg) = &input.inference_config {
        apply_inference_config(&mut body, cfg);
    }

    // 5. toolConfig
    if let Some(tc) = &input.tool_config {
        apply_tool_config(&mut body, tc);
    }

    // 7. stream flag
    if stream {
        body.insert("stream".into(), json!(true));
    }

    // 6. additionalModelRequestFields — shallow merge last (may override)
    if let Some(Value::Object(extra_map)) = &input.additional_model_request_fields {
        for (k, v) in extra_map {
            body.insert(k.clone(), v.clone());
        }
    }

    Value::Object(body)
}

/// Attach the ephemeral cache_control marker to the last emitted content part
/// (spec §3 cache points). A cachePoint with no preceding part is silently ignored.
fn apply_cache_point(parts: &mut [Value]) {
    if let Some(Value::Object(last)) = parts.last_mut() {
        last.insert("cache_control".into(), json!({ "type": "ephemeral" }));
    }
}

fn build_messages(input: &ConverseInput) -> Vec<Value> {
    let mut out: Vec<Value> = Vec::new();

    // System prompt
    if !input.system.is_empty() {
        if input.system.iter().any(|s| s.cache_point.is_some()) {
            // Cache-point mode: content becomes a parts array (spec §3).
            let mut sys_parts: Vec<Value> = Vec::new();
            for s in &input.system {
                if s.cache_point.is_some() {
                    apply_cache_point(&mut sys_parts);
                } else {
                    sys_parts.push(json!({ "type": "text", "text": s.text }));
                }
            }
            if !sys_parts.is_empty() {
                out.push(json!({ "role": "system", "content": sys_parts }));
            }
        } else {
            let text = input
                .system
                .iter()
                .map(|s| s.text.as_str())
                .collect::<Vec<_>>()
                .join("\n\n");
            out.push(json!({ "role": "system", "content": text }));
        }
    }

    for msg in &input.messages {
        match msg.role {
            ConversationRole::User => {
                // tool result blocks → individual "tool" messages first
                let mut text_parts: Vec<&str> = Vec::new();
                // Ordered text/image content parts so the multimodal path
                // preserves original block order (spec §3).
                let mut user_parts: Vec<Value> = Vec::new();
                let mut has_image = false;
                let mut has_cache_point = false;
                for block in &msg.content {
                    match block {
                        ContentBlock::ToolResult(tr) => {
                            let content = tr
                                .content
                                .iter()
                                .map(|c| match c {
                                    ToolResultContent::Text(t) => t.clone(),
                                    ToolResultContent::Json(v) => v.to_string(),
                                })
                                .collect::<Vec<_>>()
                                .join("");
                            out.push(json!({
                                "role": "tool",
                                "tool_call_id": tr.tool_use_id,
                                "content": content,
                            }));
                        }
                        ContentBlock::Text(t) => {
                            text_parts.push(t.as_str());
                            user_parts.push(json!({ "type": "text", "text": t }));
                        }
                        ContentBlock::Image(img) => {
                            has_image = true;
                            user_parts.push(image_block_to_part(img));
                        }
                        ContentBlock::CachePoint(_) => {
                            has_cache_point = true;
                            apply_cache_point(&mut user_parts);
                        }
                        ContentBlock::ToolUse(_) => {} // unexpected in user turn; skip
                    }
                }
                if (has_image || has_cache_point) && !user_parts.is_empty() {
                    out.push(json!({
                        "role": "user",
                        "content": user_parts,
                    }));
                } else if !text_parts.is_empty() && !has_cache_point {
                    // Text-only turn keeps the plain-string form (spec §3).
                    out.push(json!({
                        "role": "user",
                        "content": text_parts.join("\n"),
                    }));
                }
            }
            ConversationRole::Assistant => {
                let mut text_parts: Vec<&str> = Vec::new();
                let mut assistant_parts: Vec<Value> = Vec::new();
                let mut tool_calls: Vec<Value> = Vec::new();
                let mut has_cache_point = false;
                for block in &msg.content {
                    match block {
                        ContentBlock::Text(t) => {
                            text_parts.push(t.as_str());
                            assistant_parts.push(json!({ "type": "text", "text": t }));
                        }
                        ContentBlock::ToolUse(tu) => {
                            tool_calls.push(json!({
                                "id": tu.tool_use_id,
                                "type": "function",
                                "function": {
                                    "name": tu.name,
                                    "arguments": serde_json::to_string(&tu.input)
                                        .unwrap_or_default(),
                                },
                            }));
                        }
                        ContentBlock::CachePoint(_) => {
                            has_cache_point = true;
                            apply_cache_point(&mut assistant_parts);
                        }
                        ContentBlock::Image(_) => {} // images invalid in assistant turns; ignore
                        ContentBlock::ToolResult(_) => {} // unexpected in assistant turn; skip
                    }
                }
                let mut msg_obj = Map::new();
                msg_obj.insert("role".into(), json!("assistant"));
                if has_cache_point && !assistant_parts.is_empty() {
                    msg_obj.insert("content".into(), Value::Array(assistant_parts));
                } else if text_parts.is_empty() {
                    msg_obj.insert("content".into(), Value::Null);
                } else {
                    msg_obj.insert("content".into(), json!(text_parts.join("\n")));
                }
                if !tool_calls.is_empty() {
                    msg_obj.insert("tool_calls".into(), Value::Array(tool_calls));
                }
                out.push(Value::Object(msg_obj));
            }
        }
    }

    out
}

/// Map an image block to an OpenAI `image_url` content part (spec §3).
///
/// `source.url` is passed through verbatim; `source.bytes` is base64-encoded
/// into a `data:image/<format>;base64,…` URL.
fn image_block_to_part(img: &ImageBlock) -> Value {
    let url = match &img.source {
        ImageSource::Url(u) => u.clone(),
        ImageSource::Bytes(b) => format!(
            "data:image/{};base64,{}",
            img.format.as_str(),
            BASE64_STANDARD.encode(b)
        ),
    };
    json!({ "type": "image_url", "image_url": { "url": url } })
}

fn apply_inference_config(body: &mut Map<String, Value>, cfg: &InferenceConfiguration) {
    if let Some(v) = cfg.max_tokens {
        body.insert("max_tokens".into(), json!(v));
    }
    if let Some(v) = cfg.temperature {
        body.insert("temperature".into(), json!(v));
    }
    if let Some(v) = cfg.top_p {
        body.insert("top_p".into(), json!(v));
    }
    if let Some(v) = &cfg.stop_sequences {
        body.insert("stop".into(), json!(v));
    }
}

fn apply_tool_config(body: &mut Map<String, Value>, tc: &ToolConfiguration) {
    let tools: Vec<Value> = tc
        .tools
        .iter()
        .map(|t| {
            let spec = &t.tool_spec;
            let mut func = Map::new();
            func.insert("name".into(), json!(spec.name));
            func.insert("parameters".into(), spec.input_schema.json.clone());
            if let Some(desc) = &spec.description {
                func.insert("description".into(), json!(desc));
            }
            json!({ "type": "function", "function": func })
        })
        .collect();
    body.insert("tools".into(), Value::Array(tools));

    if let Some(choice) = &tc.tool_choice {
        let wire = match choice {
            ToolChoice::Auto => json!("auto"),
            ToolChoice::Any => json!("required"),
            ToolChoice::Tool { name } => json!({
                "type": "function",
                "function": { "name": name },
            }),
        };
        body.insert("tool_choice".into(), wire);
    }
}

// ── OpenAI response → ConverseOutput ─────────────────────────────────────────

pub(crate) fn parse_response(body: &Value, latency_ms: u64) -> Result<ConverseOutput, Error> {
    // Defensive unwrap rule: no "choices" but has "data.choices"
    let body = match body.get("data") {
        Some(d) if body.get("choices").is_none() && d.get("choices").is_some() => d,
        _ => body,
    };

    let choice = body
        .get("choices")
        .and_then(|c| c.as_array())
        .and_then(|a| a.first())
        .ok_or_else(|| Error::sdk("response missing choices[0]"))?;

    let message_val = choice
        .get("message")
        .ok_or_else(|| Error::sdk("response missing choices[0].message"))?;

    let mut content_blocks: Vec<ContentBlock> = Vec::new();

    // text content
    if let Some(text) = message_val.get("content").and_then(|v| v.as_str()) {
        if !text.is_empty() {
            content_blocks.push(ContentBlock::Text(text.to_owned()));
        }
    }

    // tool_calls
    if let Some(tcs) = message_val.get("tool_calls").and_then(|v| v.as_array()) {
        for tc in tcs {
            let id = tc
                .get("id")
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_owned();
            let func = tc.get("function").unwrap_or(&Value::Null);
            let name = func
                .get("name")
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_owned();
            let args_str = func
                .get("arguments")
                .and_then(|v| v.as_str())
                .unwrap_or("{}");
            let input: Value = serde_json::from_str(args_str)
                .unwrap_or_else(|_| Value::String(args_str.to_owned()));
            content_blocks.push(ContentBlock::ToolUse(crate::types::ToolUseBlock {
                tool_use_id: id,
                name,
                input,
            }));
        }
    }

    let message = Message {
        role: ConversationRole::Assistant,
        content: content_blocks,
    };

    let finish_reason = choice
        .get("finish_reason")
        .and_then(|v| v.as_str())
        .unwrap_or("stop");
    let stop_reason = StopReason::from_finish_reason(finish_reason);

    let usage = parse_usage(body.get("usage"));

    Ok(ConverseOutput {
        output: Some(ConverseOutputEnum::Message(message)),
        stop_reason,
        usage,
        metrics: Metrics { latency_ms },
    })
}

pub(crate) fn parse_usage(v: Option<&Value>) -> TokenUsage {
    let v = v.unwrap_or(&Value::Null);
    // Cache-read tokens: prompt_tokens_details.cached_tokens, with a
    // top-level cache_read_input_tokens fallback (spec §3).
    let cache_read = v
        .get("prompt_tokens_details")
        .and_then(|d| d.get("cached_tokens"))
        .and_then(|x| x.as_u64())
        .or_else(|| {
            v.get("cache_read_input_tokens")
                .and_then(|x| x.as_u64())
        })
        .unwrap_or(0);
    TokenUsage {
        input_tokens: v.get("prompt_tokens").and_then(|x| x.as_u64()).unwrap_or(0) as u32,
        output_tokens: v
            .get("completion_tokens")
            .and_then(|x| x.as_u64())
            .unwrap_or(0) as u32,
        total_tokens: v.get("total_tokens").and_then(|x| x.as_u64()).unwrap_or(0) as u32,
        cache_read_input_tokens: cache_read as u32,
        cache_write_input_tokens: v
            .get("cache_creation_input_tokens")
            .and_then(|x| x.as_u64())
            .unwrap_or(0) as u32,
    }
}

// ── Stream state machine ──────────────────────────────────────────────────────

/// Persistent state across SSE chunks for the streaming state machine (§3).
pub(crate) struct StreamState {
    pub message_started: bool,
    /// `messageStop` already emitted — a duplicate `finish_reason` is ignored.
    pub message_stopped: bool,
    /// Maps tool-call wire index → assigned contentBlockIndex
    pub tool_call_index_map: std::collections::HashMap<u64, u32>,
    pub text_block_index: Option<u32>,
    pub next_index: u32,
    /// Track which block indices are open (not yet stopped)
    pub open_blocks: Vec<u32>,
    /// Pending usage chunk (emitted last as metadata)
    pub pending_usage: Option<Value>,
    pub start_time: std::time::Instant,
}

impl StreamState {
    pub fn new() -> Self {
        Self {
            message_started: false,
            message_stopped: false,
            tool_call_index_map: std::collections::HashMap::new(),
            text_block_index: None,
            next_index: 0,
            open_blocks: Vec::new(),
            pending_usage: None,
            start_time: std::time::Instant::now(),
        }
    }
}

// ── Typed view of one OpenAI-style SSE chunk ──────────────────────────────────

#[derive(Deserialize)]
struct StreamChunk {
    #[serde(default)]
    choices: Vec<ChunkChoice>,
    #[serde(default)]
    usage: Option<Value>,
    #[serde(default)]
    error: Option<Value>,
}

#[derive(Deserialize)]
struct ChunkChoice {
    #[serde(default)]
    delta: Option<ChunkDelta>,
    #[serde(default)]
    finish_reason: Option<String>,
}

#[derive(Deserialize, Default)]
struct ChunkDelta {
    #[serde(default)]
    content: Option<String>,
    #[serde(default)]
    tool_calls: Vec<ChunkToolCall>,
}

#[derive(Deserialize)]
struct ChunkToolCall {
    #[serde(default)]
    index: Option<u64>,
    #[serde(default)]
    id: Option<String>,
    #[serde(default)]
    function: Option<ChunkToolFunction>,
}

#[derive(Deserialize, Default)]
struct ChunkToolFunction {
    #[serde(default)]
    name: Option<String>,
    #[serde(default)]
    arguments: Option<String>,
}

/// Process one raw SSE data payload through the state machine.
/// Returns a list of events to emit (may be empty).
pub(crate) fn process_chunk(
    state: &mut StreamState,
    raw: &str,
) -> Result<Vec<ConverseStreamOutput>, Error> {
    let chunk: StreamChunk =
        serde_json::from_str(raw).map_err(|e| Error::sdk(format!("SSE JSON parse error: {e}")))?;

    // 1. Top-level error — parse via the canonical shape-A/shape-B extractor and
    //    code→variant table (spec §4); codeless errors fall back to
    //    ModelStreamError.
    if chunk.error.is_some() {
        let (message, code, request_id, error_details) = extract_error_fields(raw);
        let details = Box::new(ErrorDetails {
            message,
            status_code: 0,
            error_code: code.clone(),
            request_id,
            details: error_details,
            retry_after_seconds: None,
            raw_body: raw.to_owned(),
        });
        return Err(match code.as_deref() {
            Some(c) => classify_by_code(c, 0, details).unwrap_or_else(Error::ModelStreamError),
            None => Error::ModelStreamError(details),
        });
    }

    let mut events: Vec<ConverseStreamOutput> = Vec::new();

    let choice = chunk.choices.first();
    let has_delta = choice.is_some_and(|c| c.delta.is_some());
    let finish_reason = choice.and_then(|c| c.finish_reason.as_deref());

    // 2. messageStart on the first chunk carrying a delta (even an empty one).
    //    A delta-less finish chunk also triggers it: per spec §3 step 5 a
    //    messageStop must always be preceded by messageStart.
    if !state.message_started && (has_delta || finish_reason.is_some()) {
        events.push(ConverseStreamOutput::MessageStart(MessageStartEvent {
            role: ConversationRole::Assistant,
        }));
        state.message_started = true;
    }

    if let Some(delta) = choice.and_then(|c| c.delta.as_ref()) {
        // 3. text delta
        if let Some(text) = delta.content.as_deref() {
            if !text.is_empty() {
                let idx = if let Some(i) = state.text_block_index {
                    i
                } else {
                    let i = state.next_index;
                    state.next_index += 1;
                    state.text_block_index = Some(i);
                    state.open_blocks.push(i);
                    i
                };
                events.push(ConverseStreamOutput::ContentBlockDelta(
                    ContentBlockDeltaEvent {
                        content_block_index: idx,
                        delta: ContentBlockDeltaPayload::Text(text.to_owned()),
                    },
                ));
            }
        }

        // 4. tool_calls deltas
        for (j, tc) in delta.tool_calls.iter().enumerate() {
            let wire_idx = tc.index.unwrap_or(j as u64);

            let block_idx = if let Some(&i) = state.tool_call_index_map.get(&wire_idx) {
                i
            } else {
                let i = state.next_index;
                state.next_index += 1;
                state.tool_call_index_map.insert(wire_idx, i);
                state.open_blocks.push(i);

                events.push(ConverseStreamOutput::ContentBlockStart(
                    ContentBlockStartEvent {
                        content_block_index: i,
                        start: ContentBlockStartPayload::ToolUse(ContentBlockStartToolUse {
                            tool_use_id: tc.id.clone().unwrap_or_default(),
                            name: tc
                                .function
                                .as_ref()
                                .and_then(|f| f.name.clone())
                                .unwrap_or_default(),
                        }),
                    },
                ));
                i
            };

            if let Some(frag) = tc.function.as_ref().and_then(|f| f.arguments.as_deref()) {
                if !frag.is_empty() {
                    events.push(ConverseStreamOutput::ContentBlockDelta(
                        ContentBlockDeltaEvent {
                            content_block_index: block_idx,
                            delta: ContentBlockDeltaPayload::ToolUse {
                                input: frag.to_owned(),
                            },
                        },
                    ));
                }
            }
        }

    }

    // 5. finish_reason — processed unconditionally (spec §3 step 5), even when
    //    the chunk carries no delta at all; messageStop is emitted at most once.
    if let Some(fr) = finish_reason {
        if !state.message_stopped {
            state.message_stopped = true;
            let stop_reason = StopReason::from_finish_reason(fr);
            // emit contentBlockStop for all open blocks in ascending order
            let mut open = std::mem::take(&mut state.open_blocks);
            open.sort_unstable();
            for idx in open {
                events.push(ConverseStreamOutput::ContentBlockStop(
                    ContentBlockStopEvent {
                        content_block_index: idx,
                    },
                ));
            }
            events.push(ConverseStreamOutput::MessageStop(MessageStopEvent {
                stop_reason,
            }));
        }
    }

    // 6. usage chunk (may arrive with empty choices)
    if let Some(usage) = chunk.usage {
        state.pending_usage = Some(usage);
    }

    Ok(events)
}

/// Build the final metadata event (emitted on stream end / [DONE]).
pub(crate) fn flush_metadata(state: &StreamState) -> Option<ConverseStreamOutput> {
    state.pending_usage.as_ref().map(|u| {
        ConverseStreamOutput::Metadata(MetadataEvent {
            usage: parse_usage(Some(u)),
            metrics: Metrics {
                latency_ms: state.start_time.elapsed().as_millis() as u64,
            },
        })
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{
        ContentBlock, ConversationRole, ImageBlock, ImageFormat, ImageSource,
        InferenceConfiguration, Message, SystemContentBlock, ToolChoice, ToolConfiguration,
        ToolInputSchema, ToolResultBlock, ToolResultContent, ToolSpecification, ToolUseBlock,
    };
    use crate::ConverseInput;
    use serde_json::json;

    fn simple_input(text: &str) -> ConverseInput {
        ConverseInput {
            model_id: "gpt-5.4-mini".into(),
            messages: vec![Message {
                role: ConversationRole::User,
                content: vec![ContentBlock::Text(text.into())],
            }],
            system: vec![],
            inference_config: None,
            tool_config: None,
            additional_model_request_fields: None,
        }
    }

    #[test]
    fn basic_request_body() {
        let input = simple_input("Hello");
        let body = build_request_body(&input, false);
        assert_eq!(body["model"], "gpt-5.4-mini");
        // no system in this input — first message is user
        assert_eq!(body["messages"][0]["role"], "user");
        assert_eq!(body["messages"][0]["content"], "Hello");
        // stream flag is absent
        assert!(body.get("stream").is_none());
    }

    #[test]
    fn system_prompt_prepended() {
        let mut input = simple_input("Q");
        input.system = vec![
            SystemContentBlock::text("You are A"),
            SystemContentBlock::text("You are B"),
        ];
        let body = build_request_body(&input, false);
        assert_eq!(body["messages"][0]["role"], "system");
        assert_eq!(body["messages"][0]["content"], "You are A\n\nYou are B");
        assert_eq!(body["messages"][1]["role"], "user");
    }

    #[test]
    fn inference_config_mapping() {
        let mut input = simple_input("Q");
        input.inference_config = Some(
            InferenceConfiguration::builder()
                .max_tokens(500)
                .temperature(0.2)
                .top_p(0.9)
                .stop_sequences(vec!["END".into()])
                .build(),
        );
        let body = build_request_body(&input, false);
        assert_eq!(body["max_tokens"], 500);
        assert_eq!(body["temperature"], 0.2);
        assert_eq!(body["top_p"], 0.9);
        assert_eq!(body["stop"], json!(["END"]));
    }

    #[test]
    fn tool_result_user_turn() {
        let input = ConverseInput {
            model_id: "m".into(),
            messages: vec![Message {
                role: ConversationRole::User,
                content: vec![
                    ContentBlock::ToolResult(ToolResultBlock {
                        tool_use_id: "call_1".into(),
                        content: vec![ToolResultContent::Text("42 degrees".into())],
                        status: None,
                    }),
                    ContentBlock::Text("Thanks".into()),
                ],
            }],
            system: vec![],
            inference_config: None,
            tool_config: None,
            additional_model_request_fields: None,
        };
        let body = build_request_body(&input, false);
        let msgs = body["messages"].as_array().unwrap();
        // tool message first
        assert_eq!(msgs[0]["role"], "tool");
        assert_eq!(msgs[0]["tool_call_id"], "call_1");
        assert_eq!(msgs[0]["content"], "42 degrees");
        // text message after
        assert_eq!(msgs[1]["role"], "user");
        assert_eq!(msgs[1]["content"], "Thanks");
    }

    fn user_input(blocks: Vec<ContentBlock>) -> ConverseInput {
        ConverseInput {
            model_id: "gpt-5.5".into(),
            messages: vec![Message {
                role: ConversationRole::User,
                content: blocks,
            }],
            system: vec![],
            inference_config: None,
            tool_config: None,
            additional_model_request_fields: None,
        }
    }

    fn image_block(format: ImageFormat, source: ImageSource) -> ContentBlock {
        ContentBlock::Image(ImageBlock { format, source })
    }

    #[test]
    fn image_bytes_to_data_url_part() {
        let input = user_input(vec![
            ContentBlock::Text("What is in this image?".into()),
            image_block(ImageFormat::Png, ImageSource::Bytes(b"hello".to_vec())),
        ]);
        let body = build_request_body(&input, false);
        let content = body["messages"][0]["content"].as_array().unwrap();
        assert_eq!(
            content[0],
            json!({ "type": "text", "text": "What is in this image?" })
        );
        assert_eq!(content[1]["type"], "image_url");
        // "hello" → "aGVsbG8=" (standard base64, known value)
        assert_eq!(
            content[1]["image_url"]["url"],
            "data:image/png;base64,aGVsbG8="
        );
    }

    #[test]
    fn image_jpeg_media_type() {
        let input = user_input(vec![image_block(
            ImageFormat::Jpeg,
            ImageSource::Bytes(b"jpegdata".to_vec()),
        )]);
        let body = build_request_body(&input, false);
        let url = body["messages"][0]["content"][0]["image_url"]["url"]
            .as_str()
            .unwrap();
        assert!(url.starts_with("data:image/jpeg;base64,"));
    }

    #[test]
    fn image_url_source_passed_verbatim() {
        let https_url = "https://example.com/cat.png";
        let input = user_input(vec![
            ContentBlock::Text("Describe".into()),
            image_block(ImageFormat::Png, ImageSource::Url(https_url.into())),
        ]);
        let body = build_request_body(&input, false);
        let content = body["messages"][0]["content"].as_array().unwrap();
        assert_eq!(content[1]["image_url"]["url"], https_url);
    }

    #[test]
    fn image_block_order_preserved() {
        let input = user_input(vec![
            image_block(ImageFormat::Png, ImageSource::Bytes(b"a".to_vec())),
            ContentBlock::Text("middle".into()),
            image_block(ImageFormat::Gif, ImageSource::Bytes(b"b".to_vec())),
        ]);
        let body = build_request_body(&input, false);
        let content = body["messages"][0]["content"].as_array().unwrap();
        let types: Vec<&str> = content
            .iter()
            .map(|p| p["type"].as_str().unwrap())
            .collect();
        assert_eq!(types, ["image_url", "text", "image_url"]);
        assert_eq!(content[1]["text"], "middle");
    }

    #[test]
    fn text_only_turn_stays_plain_string() {
        let input = user_input(vec![
            ContentBlock::Text("a".into()),
            ContentBlock::Text("b".into()),
        ]);
        let body = build_request_body(&input, false);
        assert_eq!(body["messages"][0]["content"], "a\nb");
    }

    #[test]
    fn image_with_tool_result_ordering() {
        // toolResult still becomes its own role:tool message before the user turn
        let input = user_input(vec![
            ContentBlock::ToolResult(ToolResultBlock {
                tool_use_id: "call_1".into(),
                content: vec![ToolResultContent::Text("42".into())],
                status: None,
            }),
            ContentBlock::Text("And this image?".into()),
            image_block(ImageFormat::Png, ImageSource::Bytes(b"x".to_vec())),
        ]);
        let body = build_request_body(&input, false);
        let msgs = body["messages"].as_array().unwrap();
        assert_eq!(msgs[0]["role"], "tool");
        assert_eq!(msgs[0]["tool_call_id"], "call_1");
        assert_eq!(msgs[1]["role"], "user");
        let types: Vec<&str> = msgs[1]["content"]
            .as_array()
            .unwrap()
            .iter()
            .map(|p| p["type"].as_str().unwrap())
            .collect();
        assert_eq!(types, ["text", "image_url"]);
    }

    #[test]
    fn image_in_assistant_turn_ignored() {
        let input = ConverseInput {
            model_id: "gpt-5.5".into(),
            messages: vec![
                Message {
                    role: ConversationRole::User,
                    content: vec![ContentBlock::Text("Hi".into())],
                },
                Message {
                    role: ConversationRole::Assistant,
                    content: vec![
                        ContentBlock::Text("Sure".into()),
                        image_block(ImageFormat::Png, ImageSource::Bytes(b"x".to_vec())),
                    ],
                },
                Message {
                    role: ConversationRole::User,
                    content: vec![ContentBlock::Text("Go on".into())],
                },
            ],
            system: vec![],
            inference_config: None,
            tool_config: None,
            additional_model_request_fields: None,
        };
        let body = build_request_body(&input, false);
        let msgs = body["messages"].as_array().unwrap();
        assert_eq!(msgs.len(), 3);
        assert_eq!(msgs[1]["role"], "assistant");
        // assistant content stays a plain string; the image block is dropped
        assert_eq!(msgs[1]["content"], "Sure");
        assert!(body.to_string().matches("image_url").count() == 0);
    }

    #[test]
    fn assistant_tool_calls_serialized() {
        let input = ConverseInput {
            model_id: "m".into(),
            messages: vec![Message {
                role: ConversationRole::Assistant,
                content: vec![ContentBlock::ToolUse(ToolUseBlock {
                    tool_use_id: "call_1".into(),
                    name: "get_weather".into(),
                    input: json!({"location": "NYC"}),
                })],
            }],
            system: vec![],
            inference_config: None,
            tool_config: None,
            additional_model_request_fields: None,
        };
        let body = build_request_body(&input, false);
        let msgs = body["messages"].as_array().unwrap();
        assert_eq!(msgs[0]["role"], "assistant");
        let tcs = msgs[0]["tool_calls"].as_array().unwrap();
        assert_eq!(tcs[0]["id"], "call_1");
        assert_eq!(tcs[0]["function"]["name"], "get_weather");
        let args: serde_json::Value =
            serde_json::from_str(tcs[0]["function"]["arguments"].as_str().unwrap()).unwrap();
        assert_eq!(args["location"], "NYC");
    }

    #[test]
    fn tool_choice_mapping() {
        let make = |choice: ToolChoice| {
            let mut input = simple_input("Q");
            input.tool_config = Some(ToolConfiguration {
                tools: vec![crate::types::Tool {
                    tool_spec: ToolSpecification {
                        name: "f".into(),
                        description: None,
                        input_schema: ToolInputSchema { json: json!({}) },
                    },
                }],
                tool_choice: Some(choice),
            });
            build_request_body(&input, false)
        };

        assert_eq!(make(ToolChoice::Auto)["tool_choice"], "auto");
        assert_eq!(make(ToolChoice::Any)["tool_choice"], "required");
        assert_eq!(
            make(ToolChoice::Tool { name: "f".into() })["tool_choice"]["type"],
            "function"
        );
    }

    #[test]
    fn additional_fields_merge_last() {
        let mut input = simple_input("Q");
        input.inference_config = Some(InferenceConfiguration::builder().max_tokens(100).build());
        input.additional_model_request_fields =
            Some(json!({ "max_tokens": 9999, "custom_field": true }));
        let body = build_request_body(&input, false);
        // additional fields override
        assert_eq!(body["max_tokens"], 9999);
        assert_eq!(body["custom_field"], true);
    }

    #[test]
    fn parse_response_text() {
        let raw = json!({
            "choices": [{ "message": { "role": "assistant", "content": "Hello!" }, "finish_reason": "stop" }],
            "usage": { "prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8 }
        });
        let out = parse_response(&raw, 100).unwrap();
        let msg = out.output().unwrap().as_message().unwrap();
        assert_eq!(msg.content[0].as_text().unwrap(), "Hello!");
        assert_eq!(*out.stop_reason(), StopReason::EndTurn);
        assert_eq!(out.usage().input_tokens, 5);
        assert_eq!(out.metrics().latency_ms, 100);
    }

    #[test]
    fn parse_response_tool_calls() {
        let raw = json!({
            "choices": [{ "message": {
                "role": "assistant",
                "content": null,
                "tool_calls": [{ "id": "call_1", "type": "function", "function": { "name": "get_weather", "arguments": "{\"location\":\"NYC\"}" } }]
            }, "finish_reason": "tool_calls" }],
            "usage": {}
        });
        let out = parse_response(&raw, 0).unwrap();
        assert_eq!(*out.stop_reason(), StopReason::ToolUse);
        let msg = out.output().unwrap().as_message().unwrap();
        let tu = msg.content[0].as_tool_use().unwrap();
        assert_eq!(tu.name, "get_weather");
        assert_eq!(tu.input["location"], "NYC");
    }

    #[test]
    fn defensive_unwrap_rule() {
        let raw = json!({
            "data": {
                "choices": [{ "message": { "content": "hi" }, "finish_reason": "stop" }],
                "usage": {}
            }
        });
        let out = parse_response(&raw, 0).unwrap();
        assert!(out.output().is_some());
    }

    #[test]
    fn stream_state_machine_text() {
        let mut state = StreamState::new();

        // first chunk: role delta
        let c1 = json!({"choices": [{"delta": {"role": "assistant"}, "index": 0}]});
        let evts = process_chunk(&mut state, &c1.to_string()).unwrap();
        assert!(matches!(evts[0], ConverseStreamOutput::MessageStart(_)));
        assert!(state.message_started);

        // text delta
        let c2 = json!({"choices": [{"delta": {"content": "Hello"}, "index": 0}]});
        let evts = process_chunk(&mut state, &c2.to_string()).unwrap();
        assert!(matches!(
            evts[0],
            ConverseStreamOutput::ContentBlockDelta(_)
        ));

        // finish
        let c3 = json!({"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}]});
        let evts = process_chunk(&mut state, &c3.to_string()).unwrap();
        assert!(evts
            .iter()
            .any(|e| matches!(e, ConverseStreamOutput::ContentBlockStop(_))));
        assert!(evts
            .iter()
            .any(|e| matches!(e, ConverseStreamOutput::MessageStop(_))));

        // usage chunk
        let c4 = json!({"choices": [], "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}});
        process_chunk(&mut state, &c4.to_string()).unwrap();
        let meta = flush_metadata(&state).unwrap();
        assert!(matches!(meta, ConverseStreamOutput::Metadata(_)));
    }

    #[test]
    fn stream_tool_call() {
        let mut state = StreamState::new();

        // start chunk with role
        let c1 = json!({"choices": [{"delta": {"role": "assistant"}}]});
        process_chunk(&mut state, &c1.to_string()).unwrap();

        // tool call start
        let c2 = json!({"choices": [{"delta": {"tool_calls": [{"index": 0, "id": "call_x", "type": "function", "function": {"name": "get_weather", "arguments": ""}}]}}]});
        let evts = process_chunk(&mut state, &c2.to_string()).unwrap();
        assert!(evts
            .iter()
            .any(|e| matches!(e, ConverseStreamOutput::ContentBlockStart(_))));

        // argument fragment
        let c3 = json!({"choices": [{"delta": {"tool_calls": [{"index": 0, "function": {"arguments": "{\"x\":"}}]}}]});
        let evts = process_chunk(&mut state, &c3.to_string()).unwrap();
        assert!(evts
            .iter()
            .any(|e| matches!(e, ConverseStreamOutput::ContentBlockDelta(_))));
    }

    #[test]
    fn delta_less_finish_reason_still_stops_message() {
        // Spec §3 step 5: finish_reason is processed even when the chunk has
        // no `delta` key at all (or `"delta": null`) — and a messageStop must
        // be preceded by a messageStart.
        let mut state = StreamState::new();
        let c = json!({"choices": [{"finish_reason": "stop"}]});
        let evts = process_chunk(&mut state, &c.to_string()).unwrap();
        assert!(matches!(evts[0], ConverseStreamOutput::MessageStart(_)));
        assert!(matches!(
            evts.last().unwrap(),
            ConverseStreamOutput::MessageStop(_)
        ));
    }

    #[test]
    fn delta_null_finish_reason_still_stops_message() {
        let mut state = StreamState::new();
        // start + some text first
        let c1 = json!({"choices": [{"delta": {"content": "hi"}}]});
        process_chunk(&mut state, &c1.to_string()).unwrap();

        let c2 = json!({"choices": [{"delta": null, "finish_reason": "stop"}]});
        let evts = process_chunk(&mut state, &c2.to_string()).unwrap();
        assert!(evts
            .iter()
            .any(|e| matches!(e, ConverseStreamOutput::ContentBlockStop(_))));
        assert!(evts
            .iter()
            .any(|e| matches!(e, ConverseStreamOutput::MessageStop(_))));
    }

    #[test]
    fn duplicate_finish_reason_emits_single_message_stop() {
        let mut state = StreamState::new();
        let c1 = json!({"choices": [{"delta": {"content": "hi"}}]});
        process_chunk(&mut state, &c1.to_string()).unwrap();

        let stop = json!({"choices": [{"delta": {}, "finish_reason": "stop"}]});
        let evts1 = process_chunk(&mut state, &stop.to_string()).unwrap();
        assert!(evts1
            .iter()
            .any(|e| matches!(e, ConverseStreamOutput::MessageStop(_))));

        // A second finish chunk must be ignored.
        let evts2 = process_chunk(&mut state, &stop.to_string()).unwrap();
        assert!(!evts2
            .iter()
            .any(|e| matches!(e, ConverseStreamOutput::MessageStop(_))));
        assert!(!evts2
            .iter()
            .any(|e| matches!(e, ConverseStreamOutput::ContentBlockStop(_))));
    }

    #[test]
    fn stream_error_shape_b_captures_request_id() {
        let mut state = StreamState::new();
        let raw = r#"{"success":false,"error":{"code":"STREAM_ERROR","message":"upstream broke","request_id":"req-stream-1","details":{}}}"#;
        let err = process_chunk(&mut state, raw).unwrap_err();
        match err {
            Error::ModelStreamError(d) => {
                assert_eq!(d.message, "upstream broke");
                assert_eq!(d.error_code.as_deref(), Some("STREAM_ERROR"));
                // spec §4: request_id must be captured from shape-B payloads,
                // including mid-stream error chunks
                assert_eq!(d.request_id.as_deref(), Some("req-stream-1"));
            }
            other => panic!("expected ModelStreamError, got {other:?}"),
        }
    }

    // ── cache points (spec §5 item 12) ────────────────────────────────────────

    fn cache_point() -> ContentBlock {
        ContentBlock::CachePoint(crate::types::CachePointBlock::new())
    }

    #[test]
    fn cache_point_marks_preceding_part() {
        let input = user_input(vec![
            ContentBlock::Text("Big stable document.".into()),
            cache_point(),
            ContentBlock::Text("Question?".into()),
        ]);
        let body = build_request_body(&input, false);
        let parts = body["messages"][0]["content"].as_array().unwrap();
        assert_eq!(parts.len(), 2);
        assert_eq!(parts[0]["type"], "text");
        assert_eq!(parts[0]["text"], "Big stable document.");
        assert_eq!(parts[0]["cache_control"], json!({ "type": "ephemeral" }));
        assert!(parts[1].get("cache_control").is_none());
        // cachePoint blocks never appear on the wire
        assert!(!body.to_string().contains("cachePoint"));
    }

    #[test]
    fn multiple_cache_points_mark_each_preceding_part() {
        let input = user_input(vec![
            ContentBlock::Text("part one".into()),
            cache_point(),
            ContentBlock::Text("part two".into()),
            cache_point(),
        ]);
        let body = build_request_body(&input, false);
        let parts = body["messages"][0]["content"].as_array().unwrap();
        assert_eq!(parts.len(), 2);
        assert_eq!(parts[0]["cache_control"], json!({ "type": "ephemeral" }));
        assert_eq!(parts[1]["cache_control"], json!({ "type": "ephemeral" }));
    }

    #[test]
    fn leading_cache_point_silently_ignored() {
        let input = user_input(vec![cache_point(), ContentBlock::Text("Hi".into())]);
        let body = build_request_body(&input, false);
        let parts = body["messages"][0]["content"].as_array().unwrap();
        assert_eq!(parts.len(), 1);
        assert!(parts[0].get("cache_control").is_none());
    }

    #[test]
    fn cache_point_only_turn_omitted() {
        let input = ConverseInput {
            model_id: "claude-sonnet".into(),
            messages: vec![
                Message {
                    role: ConversationRole::User,
                    content: vec![cache_point()],
                },
                Message {
                    role: ConversationRole::User,
                    content: vec![ContentBlock::Text("Hi".into())],
                },
            ],
            system: vec![],
            inference_config: None,
            tool_config: None,
            additional_model_request_fields: None,
        };
        let body = build_request_body(&input, false);
        let msgs = body["messages"].as_array().unwrap();
        assert_eq!(msgs.len(), 1);
        assert_eq!(msgs[0]["content"], "Hi");
    }

    #[test]
    fn cache_point_in_system() {
        let mut input = simple_input("Hi");
        input.system = vec![
            SystemContentBlock::text("You are a helpful assistant."),
            SystemContentBlock::cache_point(),
        ];
        let body = build_request_body(&input, false);
        assert_eq!(body["messages"][0]["role"], "system");
        let parts = body["messages"][0]["content"].as_array().unwrap();
        assert_eq!(parts[0]["cache_control"], json!({ "type": "ephemeral" }));
    }

    #[test]
    fn system_without_cache_point_unchanged() {
        let mut input = simple_input("Hi");
        input.system = vec![
            SystemContentBlock::text("A"),
            SystemContentBlock::text("B"),
        ];
        let body = build_request_body(&input, false);
        assert_eq!(body["messages"][0]["content"], "A\n\nB");
    }

    #[test]
    fn cache_point_marks_image_part() {
        let input = user_input(vec![
            image_block(ImageFormat::Png, ImageSource::Bytes(b"img".to_vec())),
            cache_point(),
            ContentBlock::Text("What is this?".into()),
        ]);
        let body = build_request_body(&input, false);
        let parts = body["messages"][0]["content"].as_array().unwrap();
        assert_eq!(parts[0]["type"], "image_url");
        assert_eq!(parts[0]["cache_control"], json!({ "type": "ephemeral" }));
        assert!(parts[1].get("cache_control").is_none());
    }

    #[test]
    fn cache_point_after_tool_result_ignored() {
        let input = user_input(vec![
            ContentBlock::ToolResult(ToolResultBlock {
                tool_use_id: "c1".into(),
                content: vec![ToolResultContent::Text("42".into())],
                status: None,
            }),
            cache_point(),
            ContentBlock::Text("next".into()),
        ]);
        let body = build_request_body(&input, false);
        let msgs = body["messages"].as_array().unwrap();
        assert_eq!(msgs[0]["role"], "tool");
        assert!(!msgs[0].to_string().contains("cache_control"));
        let parts = msgs[1]["content"].as_array().unwrap();
        assert!(parts[0].get("cache_control").is_none());
    }

    #[test]
    fn cache_point_in_assistant_turn() {
        let input = ConverseInput {
            model_id: "claude-sonnet".into(),
            messages: vec![
                Message {
                    role: ConversationRole::User,
                    content: vec![ContentBlock::Text("Hi".into())],
                },
                Message {
                    role: ConversationRole::Assistant,
                    content: vec![ContentBlock::Text("Long answer.".into()), cache_point()],
                },
                Message {
                    role: ConversationRole::User,
                    content: vec![ContentBlock::Text("Go on".into())],
                },
            ],
            system: vec![],
            inference_config: None,
            tool_config: None,
            additional_model_request_fields: None,
        };
        let body = build_request_body(&input, false);
        let parts = body["messages"][1]["content"].as_array().unwrap();
        assert_eq!(parts[0]["cache_control"], json!({ "type": "ephemeral" }));
    }

    #[test]
    fn usage_cache_counters_mapped() {
        let raw = json!({
            "choices": [{ "message": { "role": "assistant", "content": "ok" }, "finish_reason": "stop" }],
            "usage": {
                "prompt_tokens": 2048,
                "completion_tokens": 50,
                "total_tokens": 2098,
                "prompt_tokens_details": { "cached_tokens": 1024 },
                "cache_creation_input_tokens": 512
            }
        });
        let out = parse_response(&raw, 0).unwrap();
        assert_eq!(out.usage().cache_read_input_tokens, 1024);
        assert_eq!(out.usage().cache_write_input_tokens, 512);
    }

    #[test]
    fn usage_cache_counters_default_zero() {
        let raw = json!({
            "choices": [{ "message": { "role": "assistant", "content": "ok" }, "finish_reason": "stop" }],
            "usage": { "prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2 }
        });
        let out = parse_response(&raw, 0).unwrap();
        assert_eq!(out.usage().cache_read_input_tokens, 0);
        assert_eq!(out.usage().cache_write_input_tokens, 0);
    }

    #[test]
    fn stream_metadata_cache_counters() {
        let mut state = StreamState::new();
        let usage_chunk = json!({
            "choices": [],
            "usage": {
                "prompt_tokens": 100, "completion_tokens": 5, "total_tokens": 105,
                "prompt_tokens_details": { "cached_tokens": 90 },
                "cache_creation_input_tokens": 10
            }
        });
        process_chunk(&mut state, &usage_chunk.to_string()).unwrap();
        let meta = flush_metadata(&state).unwrap();
        match meta {
            ConverseStreamOutput::Metadata(m) => {
                assert_eq!(m.usage.cache_read_input_tokens, 90);
                assert_eq!(m.usage.cache_write_input_tokens, 10);
            }
            other => panic!("expected Metadata, got {other:?}"),
        }
    }

    #[test]
    fn stream_error_without_code_is_model_stream_error() {
        let mut state = StreamState::new();
        let raw = r#"{"error":{"message":"boom"}}"#;
        let err = process_chunk(&mut state, raw).unwrap_err();
        match err {
            Error::ModelStreamError(d) => {
                assert_eq!(d.message, "boom");
                assert!(d.error_code.is_none());
                assert!(d.request_id.is_none());
            }
            other => panic!("expected ModelStreamError, got {other:?}"),
        }
    }
}