use std::collections::HashSet;
use std::mem;

use super::*;

use crate::utils::{base64_decode, encode_uri, hex_encode, hmac_sha256, sha256, strip_think_tag};

use anyhow::{Context, Result, bail};
use aws_smithy_eventstream::frame::{DecodedFrame, MessageFrameDecoder};
use aws_smithy_eventstream::smithy::parse_response_headers;
use bytes::BytesMut;
use chrono::{DateTime, Utc};
use futures_util::StreamExt;
use indexmap::IndexMap;
use reqwest::{Client as ReqwestClient, Method, RequestBuilder};
use serde::Deserialize;
use serde_json::{Value, json};

#[derive(Debug, Clone, Deserialize)]
pub struct BedrockConfig {
    pub name: Option<String>,
    pub access_key_id: Option<String>,
    pub secret_access_key: Option<String>,
    pub region: Option<String>,
    pub session_token: Option<String>,
    #[serde(default)]
    pub models: Vec<ModelData>,
    pub prompt_cache: Option<bool>,
    pub patch: Option<RequestPatch>,
    pub extra: Option<ExtraConfig>,
}

impl BedrockClient {
    config_get_fn!(access_key_id, get_access_key_id);
    config_get_fn!(secret_access_key, get_secret_access_key);
    config_get_fn!(region, get_region);
    config_get_fn!(session_token, get_session_token);

    create_client_config!([
        ("access_key_id", "AWS Access Key ID", None, true),
        ("secret_access_key", "AWS Secret Access Key", None, true),
        ("region", "AWS Region", None, false),
    ]);

    fn chat_completions_builder(
        &self,
        client: &ReqwestClient,
        data: ChatCompletionsData,
    ) -> Result<RequestBuilder> {
        let access_key_id = self.get_access_key_id()?;
        let secret_access_key = self.get_secret_access_key()?;
        let region = self.get_region()?;
        let session_token = self.get_session_token().ok();
        let host = format!("bedrock-runtime.{region}.amazonaws.com");

        let model_name = &self.model.real_name();

        let uri = if data.stream {
            format!("/model/{model_name}/converse-stream")
        } else {
            format!("/model/{model_name}/converse")
        };

        let body = build_chat_completions_body(
            data,
            &self.model,
            self.config.prompt_cache.unwrap_or(true),
        )?;

        let mut request_data = RequestData::new("", body);
        self.patch_request_data(&mut request_data);
        let RequestData {
            url: _,
            headers,
            body,
            ..
        } = request_data;

        let builder = aws_fetch(
            client,
            &AwsCredentials {
                access_key_id,
                secret_access_key,
                region,
                session_token,
            },
            AwsRequest {
                method: Method::POST,
                host,
                service: "bedrock".into(),
                uri,
                querystring: "".into(),
                headers,
                body: body.to_string(),
            },
        )?;

        Ok(builder)
    }

    fn embeddings_builder(
        &self,
        client: &ReqwestClient,
        data: &EmbeddingsData,
    ) -> Result<RequestBuilder> {
        let access_key_id = self.get_access_key_id()?;
        let secret_access_key = self.get_secret_access_key()?;
        let region = self.get_region()?;
        let session_token = self.get_session_token().ok();
        let host = format!("bedrock-runtime.{region}.amazonaws.com");

        let uri = format!("/model/{}/invoke", self.model.real_name());

        let input_type = match data.query {
            true => "search_query",
            false => "search_document",
        };

        let body = json!({
            "texts": data.texts,
            "input_type": input_type,
        });

        let mut request_data = RequestData::new("", body);
        self.patch_request_data(&mut request_data);
        let RequestData {
            url: _,
            headers,
            body,
            ..
        } = request_data;

        let builder = aws_fetch(
            client,
            &AwsCredentials {
                access_key_id,
                secret_access_key,
                region,
                session_token,
            },
            AwsRequest {
                method: Method::POST,
                host,
                service: "bedrock".into(),
                uri,
                querystring: "".into(),
                headers,
                body: body.to_string(),
            },
        )?;

        Ok(builder)
    }
}

#[async_trait::async_trait]
impl Client for BedrockClient {
    client_common_fns!();

    async fn chat_completions_inner(
        &self,
        client: &ReqwestClient,
        data: ChatCompletionsData,
    ) -> Result<ChatCompletionsOutput> {
        let builder = self.chat_completions_builder(client, data)?;
        chat_completions(builder).await
    }

    async fn chat_completions_streaming_inner(
        &self,
        client: &ReqwestClient,
        handler: &mut SseHandler,
        data: ChatCompletionsData,
    ) -> Result<()> {
        let builder = self.chat_completions_builder(client, data)?;
        chat_completions_streaming(builder, handler).await
    }

    async fn embeddings_inner(
        &self,
        client: &ReqwestClient,
        data: &EmbeddingsData,
    ) -> Result<EmbeddingsOutput> {
        let builder = self.embeddings_builder(client, data)?;
        embeddings(builder).await
    }
}

async fn chat_completions(builder: RequestBuilder) -> Result<ChatCompletionsOutput> {
    let res = builder.send().await?;
    let status = res.status();
    let data: Value = res.json().await?;

    if !status.is_success() {
        catch_error(&data, status.as_u16())?;
    }

    debug!("non-stream-data: {data}");
    extract_chat_completions(&data)
}

async fn chat_completions_streaming(
    builder: RequestBuilder,
    handler: &mut SseHandler,
) -> Result<()> {
    let res = builder.send().await?;
    let status = res.status();
    if !status.is_success() {
        let data: Value = res.json().await?;
        catch_error(&data, status.as_u16())?;
        bail!("Invalid response data: {data}");
    }

    let mut function_name = String::new();
    let mut function_arguments = String::new();
    let mut function_id = String::new();
    let mut reasoning_state = 0;
    let mut thinking_text = String::new();
    let mut thinking_signature = String::new();

    let mut stream = res.bytes_stream();
    let mut buffer = BytesMut::new();
    let mut decoder = MessageFrameDecoder::new();
    while let Some(chunk) = stream.next().await {
        let chunk = chunk?;
        buffer.extend_from_slice(&chunk);
        while let DecodedFrame::Complete(message) = decoder.decode_frame(&mut buffer)? {
            let response_headers = parse_response_headers(&message)?;
            let message_type = response_headers.message_type.as_str();
            let smithy_type = response_headers.smithy_type.as_str();
            match (message_type, smithy_type) {
                ("event", _) => {
                    let data: Value = serde_json::from_slice(message.payload())?;
                    debug!("stream-data: {smithy_type} {data}");
                    match smithy_type {
                        "contentBlockStart" => {
                            if let Some(tool_use) = data["start"]["toolUse"].as_object()
                                && let (Some(id), Some(name)) = (
                                    json_str_from_map(tool_use, "toolUseId"),
                                    json_str_from_map(tool_use, "name"),
                                )
                            {
                                if !function_name.is_empty() {
                                    if function_arguments.is_empty() {
                                        function_arguments = String::from("{}");
                                    }
                                    let arguments: Value =
                                        function_arguments.parse().with_context(|| {
                                            format!("Tool call '{function_name}' has non-JSON arguments '{function_arguments}'")
                                        })?;
                                    handler.tool_call(ToolCall::new(
                                        function_name.clone(),
                                        arguments,
                                        Some(function_id.clone()),
                                    ))?;
                                }
                                function_arguments.clear();
                                function_name = name.into();
                                function_id = id.into();
                            }
                        }
                        "contentBlockDelta" => {
                            if let Some(text) = data["delta"]["text"].as_str() {
                                handler.text(text)?;
                            } else if let Some(text) =
                                data["delta"]["reasoningContent"]["text"].as_str()
                            {
                                reasoning_state = 1;
                                thinking_text.push_str(text);
                            } else if let Some(signature) =
                                data["delta"]["reasoningContent"]["signature"].as_str()
                            {
                                thinking_signature.push_str(signature);
                            } else if let Some(redacted_data) =
                                data["delta"]["reasoningContent"]["redactedContent"].as_str()
                            {
                                handler.thinking_block(ThinkingBlock::RedactedThinking {
                                    data: redacted_data.to_string(),
                                });
                            } else if let Some(input) = data["delta"]["toolUse"]["input"].as_str() {
                                function_arguments.push_str(input);
                            }
                        }
                        "contentBlockStop" => {
                            if reasoning_state == 1 {
                                reasoning_state = 0;
                                handler.thinking_block(ThinkingBlock::Thinking {
                                    thinking: mem::take(&mut thinking_text),
                                    signature: mem::take(&mut thinking_signature),
                                });
                            }
                            if !function_name.is_empty() {
                                if function_arguments.is_empty() {
                                    function_arguments = String::from("{}");
                                }
                                let arguments: Value = function_arguments.parse().with_context(|| {
                                    format!("Tool call '{function_name}' has non-JSON arguments '{function_arguments}'")
                                })?;
                                handler.tool_call(ToolCall::new(
                                    mem::take(&mut function_name),
                                    arguments,
                                    Some(mem::take(&mut function_id)),
                                ))?;
                                function_arguments.clear();
                            }
                        }
                        "metadata" => {
                            if let Some(usage) = bedrock_parse_usage(&data["usage"]) {
                                debug!("token-usage: {usage:?}");
                                handler.usage(usage);
                            }
                        }
                        _ => {}
                    }
                }
                ("exception", _) => {
                    let payload = base64_decode(message.payload())?;
                    let data = String::from_utf8_lossy(&payload);

                    bail!("Invalid response data: {data} (smithy_type: {smithy_type})")
                }
                _ => {
                    bail!(
                        "Unrecognized message, message_type: {message_type}, smithy_type: {smithy_type}",
                    );
                }
            }
        }
    }
    Ok(())
}

async fn embeddings(builder: RequestBuilder) -> Result<EmbeddingsOutput> {
    let res = builder.send().await?;
    let status = res.status();
    let data: Value = res.json().await?;

    if !status.is_success() {
        catch_error(&data, status.as_u16())?;
    }

    let res_body: EmbeddingsResBody =
        serde_json::from_value(data).context("Invalid embeddings data")?;
    Ok(res_body.embeddings)
}

#[derive(Deserialize)]
struct EmbeddingsResBody {
    embeddings: Vec<Vec<f32>>,
}

fn build_chat_completions_body(
    data: ChatCompletionsData,
    model: &Model,
    prompt_cache: bool,
) -> Result<Value> {
    let ChatCompletionsData {
        mut messages,
        temperature,
        top_p,
        reasoning_effort,
        functions,
        stream: _,
    } = data;

    let system_message = extract_system_message(&mut messages);

    let mut network_image_urls = vec![];

    let messages_len = messages.len();
    let messages: Vec<Value> = messages
        .into_iter()
        .enumerate()
        .flat_map(|(i, message)| {
            let Message { role, content } = message;
            match content {
                MessageContent::Text(text) if role.is_assistant() && i != messages_len - 1 => {
                    vec![json!({ "role": role, "content": [ { "text": strip_think_tag(&text) } ] })]
                }
                MessageContent::Text(text) => vec![json!({
                    "role": role,
                    "content": [
                        {
                            "text": text,
                        }
                    ],
                })],
                MessageContent::Array(list) => {
                    let content: Vec<_> = list
                        .into_iter()
                        .map(|item| match item {
                            MessageContentPart::Text { text } => {
                                json!({"text": text})
                            }
                            MessageContentPart::ImageUrl {
                                image_url: ImageUrl { url },
                            } => {
                                if let Some((mime_type, data)) = url
                                    .strip_prefix("data:")
                                    .and_then(|v| v.split_once(";base64,"))
                                {
                                    json!({
                                        "image": {
                                            "format": mime_type.replace("image/", ""),
                                            "source": {
                                                "bytes": data,
                                            }
                                        }
                                    })
                                } else {
                                    network_image_urls.push(url.clone());
                                    json!({ "url": url })
                                }
                            }
                        })
                        .collect();
                    vec![json!({
                        "role": role,
                        "content": content,
                    })]
                }
                MessageContent::ToolCalls(MessageContentToolCalls {
                    tool_results,
                    text,
                    sequence,
                }) => {
                    if !sequence {
                        let mut assistant_parts = vec![];
                        let mut user_parts = vec![];
                        for (index, tool_result) in tool_results.iter().enumerate() {
                            for block in &tool_result.thinking {
                                assistant_parts.push(converse_reasoning_block(block));
                            }
                            let round_text = if index == 0 && !text.is_empty() {
                                Some(text.as_str())
                            } else {
                                tool_result.text.as_deref()
                            };
                            if let Some(round_text) = round_text {
                                let round_text = strip_think_tag(round_text);
                                let round_text = round_text.trim();
                                if !round_text.is_empty() {
                                    assistant_parts.push(json!({
                                        "text": round_text,
                                    }))
                                }
                            }
                            assistant_parts.push(json!({
                                "toolUse": {
                                    "toolUseId": tool_result.call.id,
                                    "name": tool_result.call.name,
                                    "input": tool_result.call.arguments,
                                }
                            }));
                            user_parts.push(json!({
                                "toolResult": {
                                    "toolUseId": tool_result.call.id,
                                    "content": [bedrock_tool_result_content(&tool_result.output)]
                                }
                            }));
                        }
                        // Empty tool_results (reachable via deserialized sessions) must not emit an empty pair.
                        if assistant_parts.is_empty() {
                            vec![]
                        } else {
                            vec![
                                json!({ "role": "assistant", "content": assistant_parts }),
                                json!({ "role": "user", "content": user_parts }),
                            ]
                        }
                    } else {
                        // One pair per round: Claude can reuse tool_use IDs across API calls.
                        // A round boundary is detected by the presence of round text, but
                        // rounds where the model emitted only tool calls (no narration)
                        // carry no text marker. As a backstop, also split whenever a
                        // toolUseId would repeat within the current assistant message.
                        // Converse rejects duplicate toolUseIds in a single message.
                        let mut messages = vec![];
                        let mut assistant_parts: Vec<Value> = vec![];
                        let mut user_parts: Vec<Value> = vec![];
                        let mut chunk_ids: HashSet<&str> = HashSet::new();
                        for (index, tool_result) in tool_results.iter().enumerate() {
                            let id_collision = tool_result
                                .call
                                .id
                                .as_deref()
                                .is_some_and(|id| chunk_ids.contains(id));
                            if index > 0 && (tool_result.text.is_some() || id_collision) {
                                messages.push(
                                    json!({ "role": "assistant", "content": assistant_parts }),
                                );
                                messages.push(json!({ "role": "user", "content": user_parts }));
                                assistant_parts = vec![];
                                user_parts = vec![];
                                chunk_ids.clear();
                            }
                            if let Some(id) = tool_result.call.id.as_deref() {
                                chunk_ids.insert(id);
                            }
                            for block in &tool_result.thinking {
                                assistant_parts.push(converse_reasoning_block(block));
                            }
                            let round_text = if index == 0 && !text.is_empty() {
                                Some(text.as_str())
                            } else {
                                tool_result.text.as_deref()
                            };
                            if let Some(round_text) = round_text {
                                let round_text = strip_think_tag(round_text);
                                let round_text = round_text.trim();
                                if !round_text.is_empty() {
                                    assistant_parts.push(json!({
                                        "text": round_text,
                                    }))
                                }
                            }
                            assistant_parts.push(json!({
                                "toolUse": {
                                    "toolUseId": tool_result.call.id,
                                    "name": tool_result.call.name,
                                    "input": tool_result.call.arguments,
                                }
                            }));
                            user_parts.push(json!({
                                "toolResult": {
                                    "toolUseId": tool_result.call.id,
                                    "content": [bedrock_tool_result_content(&tool_result.output)]
                                }
                            }));
                        }
                        if !assistant_parts.is_empty() {
                            messages
                                .push(json!({ "role": "assistant", "content": assistant_parts }));
                            messages.push(json!({ "role": "user", "content": user_parts }));
                        }
                        messages
                    }
                }
            }
        })
        .collect();

    if !network_image_urls.is_empty() {
        bail!(
            "The model does not support network images: {:?}",
            network_image_urls
        );
    }

    let mut body = json!({
        "inferenceConfig": {},
        "messages": messages,
    });
    if let Some(v) = system_message {
        body["system"] = json!([
            {
                "text": v,
            }
        ])
    }

    if let Some(v) = model.max_tokens_param() {
        body["inferenceConfig"]["maxTokens"] = v.into();
    }
    if let Some(v) = temperature {
        body["inferenceConfig"]["temperature"] = v.into();
    }
    if let Some(v) = top_p {
        body["inferenceConfig"]["topP"] = v.into();
    }
    if let Some(v) = reasoning_effort {
        body["additionalModelRequestFields"] = json!({ "output_config": { "effort": v } });
    }
    if let Some(functions) = functions {
        let tools: Vec<_> = functions
            .iter()
            .map(|v| {
                json!({
                    "toolSpec": {
                        "name": v.name,
                        "description": v.description,
                        "inputSchema": {
                            "json": v.parameters,
                        },
                    }
                })
            })
            .collect();
        body["toolConfig"] = json!({
            "tools": tools,
        })
    }
    if prompt_cache && is_anthropic_model(model) {
        apply_prompt_cache_points(&mut body);
    }
    Ok(body)
}

/// Converse rejects cachePoint blocks for models that don't support them
/// (e.g. openai.gpt-*), so caching is only applied to Anthropic models.
fn is_anthropic_model(model: &Model) -> bool {
    let name = model.real_name();
    name.contains("anthropic.") || name.contains("claude")
}

/// Converse counterpart of the claude client's cache_control breakpoints:
/// a cachePoint after the tools, the system prompt, and the last two user
/// messages. See `apply_prompt_cache_breakpoints` in claude.rs for why the
/// last two.
fn apply_prompt_cache_points(body: &mut Value) {
    let cache_point = || json!({ "cachePoint": { "type": "default" } });
    if let Some(system) = body.get_mut("system").and_then(|v| v.as_array_mut()) {
        system.push(cache_point());
    }
    if let Some(tools) = body
        .get_mut("toolConfig")
        .and_then(|v| v.get_mut("tools"))
        .and_then(|v| v.as_array_mut())
    {
        tools.push(cache_point());
    }
    let Some(messages) = body.get_mut("messages").and_then(|v| v.as_array_mut()) else {
        return;
    };
    let mut remaining = 2;
    for message in messages.iter_mut().rev() {
        if remaining == 0 {
            break;
        }
        if message["role"] != "user" {
            continue;
        }
        // Converse rejects a cachePoint with no preceding block, so an empty
        // content array (a tool-call pair with no results) passes the
        // breakpoint to an earlier user message.
        if let Some(content) = message["content"].as_array_mut()
            && !content.is_empty()
        {
            content.push(cache_point());
            remaining -= 1;
        }
    }
}

fn bedrock_parse_usage(usage: &Value) -> Option<TokenUsage> {
    if !usage.is_object() {
        return None;
    }
    Some(TokenUsage {
        input_tokens: usage["inputTokens"].as_u64(),
        output_tokens: usage["outputTokens"].as_u64(),
        cache_creation_input_tokens: usage["cacheWriteInputTokens"].as_u64(),
        cache_read_input_tokens: usage["cacheReadInputTokens"].as_u64(),
    })
}

fn bedrock_tool_result_content(output: &Value) -> Value {
    match output {
        Value::Object(_) => json!({ "json": output }),
        Value::String(text) => json!({ "text": text }),
        _ => json!({ "text": output.to_string() }),
    }
}

/// Maps a ThinkingBlock (Anthropic tagged shape) to the Converse
/// reasoningContent union so signed blocks replay verbatim.
fn converse_reasoning_block(block: &ThinkingBlock) -> Value {
    match block {
        ThinkingBlock::Thinking {
            thinking,
            signature,
        } => json!({
            "reasoningContent": {
                "reasoningText": {
                    "text": thinking,
                    "signature": signature,
                }
            }
        }),
        ThinkingBlock::RedactedThinking { data } => json!({
            "reasoningContent": { "redactedContent": data }
        }),
    }
}

fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
    let mut text = String::new();
    let mut tool_calls = vec![];
    let mut thinking = vec![];
    if let Some(array) = data["output"]["message"]["content"].as_array() {
        for item in array {
            if let Some(v) = item["text"].as_str() {
                if !text.is_empty() {
                    text.push_str("\n\n");
                }
                text.push_str(v);
            } else if let Some(reasoning_text) =
                item["reasoningContent"]["reasoningText"].as_object()
            {
                if let Some(v) = json_str_from_map(reasoning_text, "text") {
                    thinking.push(ThinkingBlock::Thinking {
                        thinking: v.to_string(),
                        signature: json_str_from_map(reasoning_text, "signature")
                            .unwrap_or_default()
                            .to_string(),
                    });
                }
            } else if let Some(v) = item["reasoningContent"]["redactedContent"].as_str() {
                thinking.push(ThinkingBlock::RedactedThinking {
                    data: v.to_string(),
                });
            } else if let Some(tool_use) = item["toolUse"].as_object()
                && let (Some(id), Some(name), Some(input)) = (
                    json_str_from_map(tool_use, "toolUseId"),
                    json_str_from_map(tool_use, "name"),
                    tool_use.get("input"),
                )
            {
                tool_calls.push(ToolCall::new(
                    name.to_string(),
                    input.clone(),
                    Some(id.to_string()),
                ))
            }
        }
    }

    if text.is_empty() && tool_calls.is_empty() {
        bail!("Invalid response data: {data}");
    }

    let output = ChatCompletionsOutput {
        text,
        tool_calls,
        thinking,
        usage: bedrock_parse_usage(&data["usage"]),
    };
    Ok(output)
}

#[derive(Debug)]
struct AwsCredentials {
    access_key_id: String,
    secret_access_key: String,
    region: String,
    session_token: Option<String>,
}

#[derive(Debug)]
struct AwsRequest {
    method: Method,
    host: String,
    service: String,
    uri: String,
    querystring: String,
    headers: IndexMap<String, String>,
    body: String,
}

fn aws_fetch(
    client: &ReqwestClient,
    credentials: &AwsCredentials,
    request: AwsRequest,
) -> Result<RequestBuilder> {
    let AwsRequest {
        method,
        host,
        service,
        uri,
        querystring,
        mut headers,
        body,
    } = request;
    let region = &credentials.region;

    let endpoint = format!("https://{host}{uri}");

    let now: DateTime<Utc> = Utc::now();
    let amz_date = now.format("%Y%m%dT%H%M%SZ").to_string();
    let date_stamp = amz_date[0..8].to_string();
    headers.insert("host".into(), host.clone());
    headers.insert("x-amz-date".into(), amz_date.clone());
    if let Some(token) = credentials.session_token.clone() {
        headers.insert("x-amz-security-token".into(), token);
    }

    let canonical_headers = headers
        .iter()
        .map(|(key, value)| format!("{key}:{value}\n"))
        .collect::<Vec<_>>()
        .join("");

    let signed_headers = headers
        .iter()
        .map(|(key, _)| key.as_str())
        .collect::<Vec<_>>()
        .join(";");

    let payload_hash = sha256(&body);

    let canonical_request = format!(
        "{}\n{}\n{}\n{}\n{}\n{}",
        method,
        encode_uri(&uri),
        querystring,
        canonical_headers,
        signed_headers,
        payload_hash
    );

    let algorithm = "AWS4-HMAC-SHA256";
    let credential_scope = format!("{date_stamp}/{region}/{service}/aws4_request");
    let string_to_sign = format!(
        "{}\n{}\n{}\n{}",
        algorithm,
        amz_date,
        credential_scope,
        sha256(&canonical_request)
    );

    let signing_key = gen_signing_key(
        &credentials.secret_access_key,
        &date_stamp,
        region,
        &service,
    );
    let signature = hmac_sha256(&signing_key, &string_to_sign);
    let signature = hex_encode(&signature);

    let authorization_header = format!(
        "{} Credential={}/{}, SignedHeaders={}, Signature={}",
        algorithm, credentials.access_key_id, credential_scope, signed_headers, signature
    );

    headers.insert("authorization".into(), authorization_header);

    debug!("Request {endpoint} {body}");

    let mut request_builder = client.request(method, endpoint).body(body);

    for (key, value) in &headers {
        request_builder = request_builder.header(key, value);
    }

    Ok(request_builder)
}

fn gen_signing_key(key: &str, date_stamp: &str, region: &str, service: &str) -> Vec<u8> {
    let k_date = hmac_sha256(format!("AWS4{key}").as_bytes(), date_stamp);
    let k_region = hmac_sha256(&k_date, region);
    let k_service = hmac_sha256(&k_region, service);
    hmac_sha256(&k_service, "aws4_request")
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::function::{FunctionDeclaration, ToolCall, ToolResult};

    fn tool_result(output: Value) -> ToolResult {
        ToolResult {
            call: ToolCall::new(
                "fs_read".into(),
                json!({"path": "x"}),
                Some("tool_A".into()),
            ),
            output,
            text: None,
            thinking: vec![],
        }
    }

    fn sequence_tool_result(id: &str, text: Option<&str>) -> ToolResult {
        ToolResult {
            call: ToolCall::new("fs_read".into(), json!({"path": "x"}), Some(id.into())),
            output: json!("ok"),
            text: text.map(|t| t.to_string()),
            thinking: vec![],
        }
    }

    fn build_body_with(tool_results: Vec<ToolResult>, text: &str, sequence: bool) -> Value {
        let data = ChatCompletionsData {
            messages: vec![
                Message::new(MessageRole::User, MessageContent::Text("hello".to_string())),
                Message::new(
                    MessageRole::Assistant,
                    MessageContent::ToolCalls(MessageContentToolCalls {
                        tool_results,
                        text: text.to_string(),
                        sequence,
                    }),
                ),
            ],
            temperature: None,
            top_p: None,
            reasoning_effort: None,
            functions: None,
            stream: false,
        };
        build_chat_completions_body(data, &Model::new("bedrock", "test"), false).unwrap()
    }

    fn build_body(tool_results: Vec<ToolResult>) -> Value {
        build_body_with(tool_results, "", true)
    }

    fn tool_result_content_block(output: Value) -> Value {
        let body = build_body(vec![tool_result(output)]);
        body["messages"][2]["content"][0]["toolResult"]["content"][0].clone()
    }

    fn assert_unique_tool_use_ids_per_message(body: &Value) {
        for message in body["messages"].as_array().unwrap() {
            let Some(content) = message["content"].as_array() else {
                continue;
            };
            let mut seen = HashSet::new();
            for block in content {
                if let Some(id) = block["toolUse"]["toolUseId"].as_str() {
                    assert!(
                        seen.insert(id.to_string()),
                        "duplicate toolUseId `{id}` within a single assistant message: {message}"
                    );
                }
            }
        }
    }

    fn cache_fixture() -> ChatCompletionsData {
        ChatCompletionsData {
            messages: vec![
                Message::new(MessageRole::System, MessageContent::Text("sys".to_string())),
                Message::new(MessageRole::User, MessageContent::Text("hello".to_string())),
                Message::new(
                    MessageRole::Assistant,
                    MessageContent::Text("hi".to_string()),
                ),
                Message::new(MessageRole::User, MessageContent::Text("again".to_string())),
            ],
            temperature: None,
            top_p: None,
            reasoning_effort: None,
            functions: Some(vec![FunctionDeclaration {
                name: "fs_read".to_string(),
                description: "Read a file".to_string(),
                parameters: Default::default(),
                agent: false,
            }]),
            stream: false,
        }
    }

    #[test]
    fn test_bedrock_prompt_cache_inserts_cache_points_for_anthropic_models() {
        let body = build_chat_completions_body(
            cache_fixture(),
            &Model::new("bedrock", "anthropic.claude-sonnet-4-5-20250929-v1:0"),
            true,
        )
        .unwrap();

        let cache_point = json!({ "cachePoint": { "type": "default" } });
        assert_eq!(
            body["system"].as_array().unwrap().last(),
            Some(&cache_point),
            "body: {body}"
        );
        assert_eq!(
            body["toolConfig"]["tools"].as_array().unwrap().last(),
            Some(&cache_point),
            "body: {body}"
        );
        let messages = body["messages"].as_array().unwrap();
        for idx in [0, 2] {
            assert_eq!(
                messages[idx]["content"].as_array().unwrap().last(),
                Some(&cache_point),
                "body: {body}"
            );
        }
        assert!(
            messages[1]["content"]
                .as_array()
                .unwrap()
                .iter()
                .all(|block| block.get("cachePoint").is_none()),
            "assistant messages must not get cache points: {body}"
        );
    }

    #[test]
    fn test_bedrock_prompt_cache_skips_non_anthropic_models() {
        let body = build_chat_completions_body(
            cache_fixture(),
            &Model::new("bedrock", "openai.gpt-oss-120b-1:0"),
            true,
        )
        .unwrap();

        assert!(!body.to_string().contains("cachePoint"), "body: {body}");
    }

    /// Converse rejects a cachePoint with no preceding block, so a user
    /// message with an empty content array must be skipped and the
    /// breakpoint placed on an earlier user message instead.
    #[test]
    fn test_bedrock_prompt_cache_skips_empty_user_content() {
        let mut body = json!({
            "messages": [
                { "role": "user", "content": [ { "text": "hello" } ] },
                { "role": "assistant", "content": [ { "text": "hi" } ] },
                { "role": "user", "content": [ { "text": "again" } ] },
                { "role": "assistant", "content": [ { "text": "calling" } ] },
                { "role": "user", "content": [] },
            ],
        });
        apply_prompt_cache_points(&mut body);

        let cache_point = json!({ "cachePoint": { "type": "default" } });
        let messages = body["messages"].as_array().unwrap();
        assert!(
            messages[4]["content"].as_array().unwrap().is_empty(),
            "empty user content must not get a cachePoint: {body}"
        );
        for idx in [0, 2] {
            assert_eq!(
                messages[idx]["content"].as_array().unwrap().last(),
                Some(&cache_point),
                "body: {body}"
            );
        }
    }

    #[test]
    fn sequence_splits_on_round_text() {
        let body = build_body(vec![
            sequence_tool_result("toolu_A", None),
            sequence_tool_result("toolu_B", None),
            sequence_tool_result("toolu_C", Some("running another tool")),
        ]);

        let messages = body["messages"].as_array().unwrap();

        assert_eq!(messages.len(), 5, "body: {body}");
        assert_unique_tool_use_ids_per_message(&body);
    }

    #[test]
    fn sequence_splits_on_reused_id_in_textless_round() {
        let body = build_body(vec![
            sequence_tool_result("toolu_A", None),
            sequence_tool_result("toolu_B", None),
            sequence_tool_result("toolu_A", None),
        ]);

        let messages = body["messages"].as_array().unwrap();

        assert_eq!(messages.len(), 5, "body: {body}");
        assert_unique_tool_use_ids_per_message(&body);
    }

    #[test]
    fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() {
        let body = build_body(vec![
            sequence_tool_result("toolu_A", None),
            sequence_tool_result("toolu_B", None),
            sequence_tool_result("toolu_C", None),
        ]);

        let messages = body["messages"].as_array().unwrap();

        assert_eq!(messages.len(), 3, "body: {body}");
        assert_unique_tool_use_ids_per_message(&body);
    }

    #[test]
    fn non_sequence_emits_assistant_user_pair() {
        let body = build_body_with(
            vec![
                sequence_tool_result("toolu_A", None),
                sequence_tool_result("toolu_B", None),
            ],
            "",
            false,
        );

        let messages = body["messages"].as_array().unwrap();

        assert_eq!(messages.len(), 3, "body: {body}");
        assert_eq!(messages[1]["role"], "assistant");
        assert_eq!(messages[1]["content"].as_array().unwrap().len(), 2);
        assert_eq!(messages[2]["role"], "user");
        assert_eq!(messages[2]["content"].as_array().unwrap().len(), 2);
    }

    #[test]
    fn non_sequence_empty_tool_results_emits_no_messages() {
        let body = build_body_with(vec![], "leftover text", false);

        let messages = body["messages"].as_array().unwrap();

        assert_eq!(messages.len(), 1, "body: {body}");
    }

    #[test]
    fn sequence_empty_tool_results_emits_no_messages() {
        let body = build_body_with(vec![], "leftover text", true);

        let messages = body["messages"].as_array().unwrap();

        assert_eq!(messages.len(), 1, "body: {body}");
    }

    #[test]
    fn test_bedrock_tool_result_object_uses_json_block() {
        assert_eq!(
            tool_result_content_block(json!({"output": "ok"})),
            json!({"json": {"output": "ok"}})
        );
        assert_eq!(tool_result_content_block(json!({})), json!({"json": {}}));
    }

    #[test]
    fn test_bedrock_tool_result_array_uses_text_block() {
        assert_eq!(
            tool_result_content_block(json!([{"id": 1}, {"id": 2}])),
            json!({"text": "[{\"id\":1},{\"id\":2}]"})
        );
    }

    #[test]
    fn test_bedrock_tool_result_scalars_use_text_block() {
        assert_eq!(
            tool_result_content_block(json!("plain text")),
            json!({"text": "plain text"})
        );
        assert_eq!(tool_result_content_block(json!(42)), json!({"text": "42"}));
        assert_eq!(
            tool_result_content_block(json!(null)),
            json!({"text": "null"})
        );
    }

    fn thinking_tool_result() -> ToolResult {
        ToolResult {
            call: ToolCall::new(
                "fs_read".into(),
                json!({"path": "x"}),
                Some("tool_A".into()),
            ),
            output: json!("ok"),
            text: None,
            thinking: vec![
                ThinkingBlock::Thinking {
                    thinking: "let me think".into(),
                    signature: "sig123".into(),
                },
                ThinkingBlock::RedactedThinking {
                    data: "b64data".into(),
                },
            ],
        }
    }

    #[test]
    fn replays_thinking_blocks_in_converse_shape() {
        for sequence in [false, true] {
            let body = build_body_with(vec![thinking_tool_result()], "", sequence);

            let content = body["messages"][1]["content"].as_array().unwrap();
            assert_eq!(
                content[0],
                json!({
                    "reasoningContent": {
                        "reasoningText": { "text": "let me think", "signature": "sig123" }
                    }
                }),
                "body: {body}"
            );
            assert_eq!(
                content[1],
                json!({ "reasoningContent": { "redactedContent": "b64data" } }),
                "body: {body}"
            );
            assert!(content[2].get("toolUse").is_some(), "body: {body}");
        }
    }

    #[test]
    fn extract_chat_completions_captures_reasoning_content() {
        let data = json!({
            "output": {
                "message": {
                    "content": [
                        {
                            "reasoningContent": {
                                "reasoningText": { "text": "let me think", "signature": "sig123" }
                            }
                        },
                        { "reasoningContent": { "redactedContent": "b64data" } },
                        { "text": "answer" },
                        {
                            "toolUse": {
                                "toolUseId": "tool_A",
                                "name": "fs_read",
                                "input": { "path": "x" }
                            }
                        }
                    ]
                }
            }
        });

        let output = extract_chat_completions(&data).unwrap();

        assert_eq!(output.text, "answer");
        assert_eq!(output.tool_calls.len(), 1);
        assert_eq!(output.thinking.len(), 2);
        match &output.thinking[0] {
            ThinkingBlock::Thinking {
                thinking,
                signature,
            } => {
                assert_eq!(thinking, "let me think");
                assert_eq!(signature, "sig123");
            }
            other => panic!("unexpected block: {other:?}"),
        }
        assert!(matches!(
            &output.thinking[1],
            ThinkingBlock::RedactedThinking { data } if data == "b64data"
        ));
    }
}