module RubyLLM::Providers::Gemini::Embeddings

def embed(text, model:) # rubocop:disable Metrics/AbcSize,Metrics/MethodLength
  payload = {
    content: {
      parts: format_text_for_embedding(text)
    }
  }
  url = "models/#{model}:embedContent"
  response = post(url, payload)
  if text.is_a?(Array)
    # We need to make separate calls for each text with Gemini
    embeddings = text.map do |t|
      single_payload = { content: { parts: [{ text: t.to_s }] } }
      single_response = post(url, single_payload)
      single_response.body.dig('embedding', 'values')
    end
    Embedding.new(
      vectors: embeddings,
      model: model,
      input_tokens: response.body.dig('usageMetadata', 'promptTokenCount') || 0
    )
  else
    Embedding.new(
      vectors: response.body.dig('embedding', 'values'),
      model: model,
      input_tokens: response.body.dig('usageMetadata', 'promptTokenCount') || 0
    )
  end
end

def format_text_for_embedding(text)
  if text.is_a?(Array)
    text.map { |t| { text: t.to_s } }
  else
    [{ text: text.to_s }]
  end
end