module RubyLLM::Providers::Bedrock::Chat

# File lib/ruby_llm/providers/bedrock/chat.rb, line 320
def budget_reasoning_config(thinking)
  budget = thinking.respond_to?(:budget) ? thinking.budget : thinking
  return nil unless budget.is_a?(Integer)

  { reasoning_config: { type: 'enabled', budget_tokens: budget } }
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 279
def build_output_config(schema)
  return nil unless schema

  cleaned = RubyLLM::Utils.deep_dup(schema[:schema])
  cleaned.delete(:strict)
  cleaned.delete('strict')

  {
    textFormat: {
      type: 'json_schema',
      structure: {
        jsonSchema: {
          schema: JSON.generate(cleaned),
          name: schema[:name]
        }
      }
    }
  }
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 12
def completion_url
  "/model/#{@model.id}/converse"
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 395
def default_input_schema
  {
    'type' => 'object',
    'properties' => {},
    'required' => []
  }
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 308
def effort_reasoning_config(thinking)
  effort = thinking.respond_to?(:effort) ? thinking.effort : nil
  effort = effort.to_s if effort
  return nil if effort.nil? || effort.empty? || effort == 'none'

  if reasoning_embedded?(@model)
    { reasoning_config: { type: 'enabled', reasoning_effort: effort } }
  else
    { reasoning_effort: effort }
  end
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 71
def input_tokens(usage)
  input_tokens = usage['inputTokens']
  return unless input_tokens

  [input_tokens.to_i - usage['cacheReadInputTokens'].to_i - usage['cacheWriteInputTokens'].to_i, 0].max
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 195
def normalize_tool_result_block(block)
  return nil unless block.is_a?(Hash)
  return block if tool_result_content_block?(block)

  nil
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 48
def parse_completion_response(response)
  data = response.body
  return if data.nil? || data.empty?

  content_blocks = data.dig('output', 'message', 'content') || []
  usage = data['usage'] || {}
  thinking_text, thinking_signature = parse_thinking(content_blocks)

  Message.new(
    role: :assistant,
    content: parse_text_content(content_blocks),
    thinking: Thinking.build(text: thinking_text, signature: thinking_signature),
    tool_calls: parse_tool_calls(content_blocks),
    input_tokens: input_tokens(usage),
    output_tokens: usage['outputTokens'],
    cached_tokens: usage['cacheReadInputTokens'],
    cache_creation_tokens: usage['cacheWriteInputTokens'],
    thinking_tokens: usage['reasoningTokens'],
    model_id: data['modelId'],
    raw: response
  )
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 366
def parse_reasoning_content_block(block)
  reasoning_content = block['reasoningContent']
  return [nil, nil] unless reasoning_content.is_a?(Hash)

  reasoning_text = reasoning_content['reasoningText'] || {}
  text = reasoning_text['text'].is_a?(String) ? reasoning_text['text'] : nil
  signature = reasoning_text['signature'] if reasoning_text['signature'].is_a?(String)
  signature ||= reasoning_content['redactedContent'] if reasoning_content['redactedContent'].is_a?(String)
  [text, signature]
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 348
def parse_text_content(content_blocks)
  text = content_blocks.filter_map { |block| block['text'] if block['text'].is_a?(String) }.join
  text.empty? ? nil : text
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 353
def parse_thinking(content_blocks)
  text = +''
  signature = nil

  content_blocks.each do |block|
    chunk_text, chunk_signature = parse_reasoning_content_block(block)
    text << chunk_text if chunk_text
    signature ||= chunk_signature
  end

  [text.empty? ? nil : text, signature]
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 377
def parse_tool_calls(content_blocks)
  tool_calls = {}

  content_blocks.each do |block|
    tool_use = block['toolUse']
    next unless tool_use

    tool_call_id = tool_use['toolUseId']
    tool_calls[tool_call_id] = ToolCall.new(
      id: tool_call_id,
      name: tool_use['name'],
      arguments: tool_use['input'] || {}
    )
  end

  tool_calls.empty? ? nil : tool_calls
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 270
def render_additional_model_request_fields(thinking)
  fields = {}

  reasoning_fields = render_reasoning_fields(thinking)
  fields = RubyLLM::Utils.deep_merge(fields, reasoning_fields) if reasoning_fields

  fields.empty? ? nil : fields
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 172
def render_content_tool_result_content(content)
  blocks = []
  blocks << text_tool_result_block(content.text) unless content.text.to_s.empty?
  content.attachments.each { |attachment| blocks << text_tool_result_block(attachment.for_llm) }
  blocks.empty? ? [text_tool_result_block(nil)] : blocks
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 219
def render_inference_config(_model, temperature)
  config = {}
  config[:temperature] = temperature unless temperature.nil?
  config
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 111
def render_message_content(msg)
  if msg.content.is_a?(RubyLLM::Content::Raw)
    return render_raw_content(msg.content) if msg.role == :assistant

    return sanitize_non_assistant_raw_blocks(render_raw_content(msg.content))
  end

  blocks = []

  thinking_block = render_thinking_block(msg.thinking)
  blocks << thinking_block if msg.role == :assistant && thinking_block

  text_and_media_blocks = Media.render_content(msg.content, used_document_names: @used_document_names)
  blocks.concat(text_and_media_blocks) if text_and_media_blocks

  if msg.tool_call?
    msg.tool_calls.each_value do |tool_call|
      blocks << {
        toolUse: {
          toolUseId: tool_call.id,
          name: tool_call.name,
          input: tool_call.arguments
        }
      }
    end
  end

  blocks
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 78
def render_messages(messages)
  rendered = []
  tool_result_blocks = []

  messages.each do |msg|
    if msg.tool_result?
      tool_result_blocks << render_tool_result_block(msg)
      next
    end

    unless tool_result_blocks.empty?
      rendered << { role: 'user', content: tool_result_blocks }
      tool_result_blocks = []
    end

    message = render_non_tool_message(msg)
    rendered << message if message
  end

  rendered << { role: 'user', content: tool_result_blocks } unless tool_result_blocks.empty?
  rendered
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 101
def render_non_tool_message(msg)
  content = render_message_content(msg)
  return nil if content.empty?

  {
    role: render_role(msg.role),
    content: content
  }
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 17
def render_payload(messages, tools:, temperature:, model:, stream: false,
                   schema: nil, thinking: nil, tool_prefs: nil)
  tool_prefs ||= {}
  @model = model
  @used_document_names = {}
  system_messages, chat_messages = messages.partition { |msg| msg.role == :system }
  payload = {
    messages: render_messages(chat_messages)
  }

  system_blocks = render_system(system_messages)
  payload[:system] = system_blocks unless system_blocks.empty?

  payload[:inferenceConfig] = render_inference_config(model, temperature)

  tool_config = render_tool_config(tools, tool_prefs)
  if tool_config
    payload[:toolConfig] = tool_config
    payload[:tools] = tool_config[:tools] # Internal mirror for shared payload inspections in specs.
  end

  additional_fields = render_additional_model_request_fields(thinking)
  payload[:additionalModelRequestFields] = additional_fields if additional_fields

  output_config = build_output_config(schema)
  payload[:outputConfig] = output_config if output_config

  payload
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 141
def render_raw_content(content)
  value = content.value
  value.is_a?(Array) ? value : [value]
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 185
def render_raw_tool_result_content(raw_value)
  blocks = raw_value.is_a?(Array) ? raw_value : [raw_value]

  normalized = blocks.filter_map do |block|
    normalize_tool_result_block(block)
  end

  normalized.empty? ? [{ text: raw_value.to_s }] : normalized
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 299
def render_reasoning_fields(thinking)
  return nil unless thinking&.enabled?

  effort_config = effort_reasoning_config(thinking)
  return effort_config if effort_config

  budget_reasoning_config(thinking)
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 208
def render_role(role)
  case role
  when :assistant then 'assistant'
  else 'user'
  end
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 215
def render_system(messages)
  messages.flat_map { |msg| Media.render_content(msg.content, used_document_names: @used_document_names) }
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 327
def render_thinking_block(thinking)
  return nil unless thinking

  if thinking.text
    {
      reasoningContent: {
        reasoningText: {
          text: thinking.text,
          signature: thinking.signature
        }.compact
      }
    }
  elsif thinking.signature
    {
      reasoningContent: {
        redactedContent: thinking.signature
      }
    }
  end
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 252
def render_tool(tool)
  input_schema = tool.params_schema || RubyLLM::Tool::SchemaDefinition.from_parameters(tool.parameters)&.json_schema

  tool_spec = {
    toolSpec: {
      name: tool.name,
      description: tool.description,
      inputSchema: {
        json: input_schema || default_input_schema
      }
    }
  }

  return tool_spec if tool.provider_params.empty?

  RubyLLM::Utils.deep_merge(tool_spec, tool.provider_params)
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 239
def render_tool_choice(choice)
  case choice
  when :auto
    { auto: {} }
  when :none
    nil
  when :required
    { any: {} }
  else
    { tool: { name: choice.to_s } }
  end
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 225
def render_tool_config(tools, tool_prefs)
  return nil if tools.empty?

  config = {
    tools: tools.values.map { |tool| render_tool(tool) }
  }

  return config if tool_prefs.nil? || tool_prefs[:choice].nil?

  tool_choice = render_tool_choice(tool_prefs[:choice])
  config[:toolChoice] = tool_choice if tool_choice
  config
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 155
def render_tool_result_block(msg)
  {
    toolResult: {
      toolUseId: msg.tool_call_id,
      content: render_tool_result_content(msg.content)
    }
  }
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 164
def render_tool_result_content(content)
  return render_raw_tool_result_content(content.value) if content.is_a?(RubyLLM::Content::Raw)
  return [{ json: content }] if content.is_a?(Hash) || content.is_a?(Array)
  return render_content_tool_result_content(content) if content.is_a?(RubyLLM::Content)

  [text_tool_result_block(content)]
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 146
def sanitize_non_assistant_raw_blocks(blocks)
  blocks.filter_map do |block|
    next unless block.is_a?(Hash)
    next if block.key?(:reasoningContent) || block.key?('reasoningContent')

    block
  end
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 179
def text_tool_result_block(text)
  text = text.to_s
  text = '(no output)' if text.empty?
  { text: text }
end

# File lib/ruby_llm/providers/bedrock/chat.rb, line 202
def tool_result_content_block?(block)
  %w[text json document image].any? do |key|
    block.key?(key) || block.key?(key.to_sym)
  end
end