module RubyLLM::Providers::Gemini::Capabilities

def context_length(model_id)
  context_window_for(model_id)
end

def context_window_for(model_id)
  case model_id
  when /gemini-2\.5-pro-exp-03-25/, /gemini-2\.0-flash/, /gemini-2\.0-flash-lite/, /gemini-1\.5-flash/, /gemini-1\.5-flash-8b/ # rubocop:disable Layout/LineLength
    1_048_576
  when /gemini-1\.5-pro/ then 2_097_152
  when /gemini-embedding-exp/ then 8_192
  when /text-embedding-004/, /embedding-001/ then 2_048
  when /aqa/ then 7_168
  when /imagen-3/ then nil # No token limit for image generation
  else 32_768 # Sensible default for unknown models
  end
end

def default_input_price
  0.075 # Default to Flash pricing
end

def default_output_price
  0.30 # Default to Flash pricing
end

def format_display_name(model_id)
  model_id
    .delete_prefix('models/')
    .split('-')
    .map(&:capitalize)
    .join(' ')
    .gsub(/(\d+\.\d+)/, ' \1') # Add space before version numbers
    .gsub(/\s+/, ' ')          # Clean up multiple spaces
    .gsub('Aqa', 'AQA')        # Special case for AQA
    .strip
end

def input_price_for(model_id)
  base_price = PRICES.dig(pricing_family(model_id), :input) || default_input_price
  return base_price unless long_context_model?(model_id)
  # Apply different pricing for prompts longer than 128k tokens
  context_window_for(model_id) > 128_000 ? base_price * 2 : base_price
end

def long_context_model?(model_id)
  model_id.match?(/gemini-1\.5-(?:pro|flash)|gemini-1\.5-flash-8b/)
end

def max_tokens_for(model_id)
  case model_id
  when /gemini-2\.5-pro-exp-03-25/ then 64_000
  when /gemini-2\.0-flash/, /gemini-2\.0-flash-lite/, /gemini-1\.5-flash/, /gemini-1\.5-flash-8b/, /gemini-1\.5-pro/ # rubocop:disable Layout/LineLength
    8_192
  when /gemini-embedding-exp/ then nil # Elastic, supports 3072, 1536, or 768
  when /text-embedding-004/, /embedding-001/ then 768 # Output dimension size for embeddings
  when /aqa/ then 1_024
  when /imagen-3/ then 4 # Output images
  else 4_096 # Sensible default
  end
end

def model_family(model_id) # rubocop:disable Metrics/CyclomaticComplexity,Metrics/MethodLength
  case model_id
  when /gemini-2\.5-pro-exp-03-25/ then 'gemini25_pro_exp'
  when /gemini-2\.0-flash-lite/ then 'gemini20_flash_lite'
  when /gemini-2\.0-flash/ then 'gemini20_flash'
  when /gemini-1\.5-flash-8b/ then 'gemini15_flash_8b'
  when /gemini-1\.5-flash/ then 'gemini15_flash'
  when /gemini-1\.5-pro/ then 'gemini15_pro'
  when /gemini-embedding-exp/ then 'gemini_embedding_exp'
  when /text-embedding-004/ then 'embedding4'
  when /embedding-001/ then 'embedding1'
  when /aqa/ then 'aqa'
  when /imagen-3/ then 'imagen3'
  else 'other'
  end
end

def model_type(model_id)
  case model_id
  when /text-embedding|embedding|gemini-embedding/ then 'embedding'
  when /imagen/ then 'image'
  else 'chat'
  end
end

def output_price_for(model_id)
  base_price = PRICES.dig(pricing_family(model_id), :output) || default_output_price
  return base_price unless long_context_model?(model_id)
  # Apply different pricing for prompts longer than 128k tokens
  context_window_for(model_id) > 128_000 ? base_price * 2 : base_price
end

def pricing_family(model_id) # rubocop:disable Metrics/CyclomaticComplexity,Metrics/MethodLength
  case model_id
  when /gemini-2\.5-pro-exp-03-25/ then :pro_2_5 # rubocop:disable Naming/VariableNumber
  when /gemini-2\.0-flash-lite/ then :flash_lite_2 # rubocop:disable Naming/VariableNumber
  when /gemini-2\.0-flash/ then :flash_2 # rubocop:disable Naming/VariableNumber
  when /gemini-1\.5-flash-8b/ then :flash_8b
  when /gemini-1\.5-flash/ then :flash
  when /gemini-1\.5-pro/ then :pro
  when /gemini-embedding-exp/ then :gemini_embedding
  when /text-embedding|embedding/ then :embedding
  when /imagen/ then :imagen
  when /aqa/ then :aqa
  else :base
  end
end

def supports_audio?(model_id)
  model_id.match?(/gemini|pro|flash/)
end

def supports_caching?(model_id)
  if model_id.match?(/flash-lite|gemini-2\.5-pro-exp-03-25|aqa|imagen|text-embedding|embedding-001/)
    return false
  end
  model_id.match?(/gemini|pro|flash/)
end

def supports_functions?(model_id)
  return false if model_id.match?(/text-embedding|embedding-001|aqa|flash-lite|imagen|gemini-2\.0-flash-lite/)
  model_id.match?(/gemini|pro|flash/)
end

def supports_json_mode?(model_id)
  if model_id.match?(/text-embedding|embedding-001|aqa|imagen|gemini-2\.0-flash-lite|gemini-2\.5-pro-exp-03-25/)
    return false
  end
  model_id.match?(/gemini|pro|flash/)
end

def supports_tuning?(model_id)
  model_id.match?(/gemini-1\.5-flash|gemini-1\.5-flash-8b/)
end

def supports_vision?(model_id)
  return false if model_id.match?(/text-embedding|embedding-001|aqa/)
  model_id.match?(/gemini|flash|pro|imagen/)
end