class LittleGhost::Providers::Bedrock

# File lib/little_ghost/providers/bedrock.rb, line 70
def initialize(model:, region: nil, client: nil, max_retries: 2, sleeper: nil,
  on_retry: ->(*) {}, max_embedding_response_bytes: DEFAULT_MAX_EMBEDDING_RESPONSE_BYTES, **client_options)
  @model = model
  @client = client || build_client(region:, **client_options)
  @max_retries = Integer(max_retries)
  @sleeper = sleeper
  @on_retry = on_retry
  @max_embedding_response_bytes = positive_integer(max_embedding_response_bytes, :max_embedding_response_bytes)
end

# File lib/little_ghost/providers/bedrock.rb, line 31
def self.request_options
  %i[max_response_bytes max_retries open_timeout read_timeout].freeze
end

# File lib/little_ghost/providers/bedrock.rb, line 183
def capabilities(metadata: {})
  parameters = metadata[:supported_parameters] || metadata["supported_parameters"]
  return ModelCapabilities.unknown unless parameters.is_a?(Array)

  values = parameters.map(&:to_s)
  ModelCapabilities.new(
    native_structured_output: values.include?("structured_outputs"),
    tools: values.include?("tools"),
    tool_choice: values.include?("tool_choice"),
    supported_parameters: values
  )
end

# File lib/little_ghost/providers/bedrock.rb, line 144
def embed(request)
  unless model == "amazon.titan-embed-text-v2:0"
    raise UnsupportedModelOperationError, "Bedrock embeddings require amazon.titan-embed-text-v2:0"
  end

  dimensions = Integer(request.settings.fetch(:dimensions, 1024))
  raise ConfigurationError, "dimensions must be 256, 512, or 1024" unless [256, 512, 1024].include?(dimensions)
  normalize = request.settings.fetch(:normalize, true)
  unless normalize == true || normalize == false
    raise ConfigurationError, "normalize must be true or false"
  end

  vectors = []
  usage = Usage.new
  request.inputs.each do |input|
    response = with_retries(request) do
      @client.invoke_model(
        model_id: model,
        body: {input_text: input, dimensions:, normalize:},
        cancellation_token: request.cancellation_token,
        deadline: request.deadline,
        max_response_bytes: @max_embedding_response_bytes
      )
    end
    payload = JSON.parse(response.body)
    vector = payload.fetch("embedding")
    unless vector.is_a?(Array) && vector.length == dimensions
      raise ProtocolError, "Bedrock returned an embedding with unexpected dimensions"
    end
    vectors << vector
    usage += Usage.new(input_tokens: payload["inputTextTokenCount"])
  rescue JSON::ParserError, KeyError
    raise ProtocolError, "Bedrock returned an invalid embedding response"
  end
  Embeddings::Response.new(vectors:, usage:, metadata: {model:})
end

# File lib/little_ghost/providers/bedrock.rb, line 85
def stream(request)
  return enum_for(__method__, request) unless block_given?

  attempts = 0

  begin
    partial_text = false
    request.cancellation_token.raise_if_cancelled!
    normalizer = StreamNormalizer.new(model:)
    stream = Support::InterruptibleStream.new(
      cancellation_token: request.cancellation_token,
      deadline: request.deadline
    ) do |emit|
      response = @client.converse_stream(**request_parameters(request))
      response.stream.each { |event| emit.call(event) }
    end
    stream.each do |event|
      normalizer.consume(event_hash(event)).each do |normalized|
        partial_text ||= normalized.type == :text_delta && !normalized.data[:text].to_s.empty?
        yield normalized
      end
    end
    normalizer.finish.each do |event|
      partial_text ||= event.type == :text_delta && !event.data[:text].to_s.empty?
      yield event
    end
  rescue CancelledError, DeadlineExceededError, CleanupError
    raise
  rescue => error
    raise if error.is_a?(Error) && !error.is_a?(StreamError)

    if context_window_overflow?(error)
      raise ContextWindowOverflowError, "The model context window was exceeded"
    end
    raise provider_error(error) if !retryable?(error) || attempts >= @max_retries

    attempts += 1
    request.cancellation_token.raise_if_cancelled!
    delay = capped_retry_delay(request, retry_delay(attempts))
    @on_retry.call(attempts, error, delay)
    wait_before_retry(request, delay)
    yield StreamEvent.build(
      :model_retry,
      attempt: attempts,
      delay:,
      error_class: error.class.name,
      error_code: (error.event_type if error.is_a?(StreamError)),
      partial_text:
    )
    retry
  end
end