lib/tasks/setup_ai_gateway.rake

# frozen_string_literal: true

require_relative '../gdk'

AI_GATEWAY_ENV_FILE = File.join(GDK.root, 'gitlab-ai-gateway', '.env')
ENV_RUNIT_FILE = File.join(GDK.root, 'env.runit')
LOG_FILE = File.join(GDK.root, 'log/gitlab-ai-gateway/gateway_debug.log')
DEBUG_VARS = {
  'AIGW_LOGGING__LEVEL' => 'debug',
  'AIGW_LOGGING__FORMAT_JSON' => 'false',
  'AIGW_LOGGING__TO_FILE' => LOG_FILE
}.freeze

def update_env_file(env_file, key, value)
  env_contents = File.exist?(env_file) ? File.read(env_file).dup : ''
  env_contents = env_contents.strip

  if env_contents.match?(/^#{key}=/)
    env_contents.sub!(/^#{key}=.*/, "#{key}=#{value}")
  else
    env_contents << "\n" unless env_contents.empty?
    env_contents << "#{key}=#{value}"
  end

  File.write(env_file, "#{env_contents}\n")
end

desc 'Set up GitLab AI Gateway'
task :setup_ai_gateway do
  GDK::Output.puts 'Setting up GitLab AI Gateway...'
  anthropic_key = GDK::Output.prompt('Enter your Anthropic API key')
  enable_debug = GDK::Output.prompt('Do you want to set additional environment variables for debugging? [y/N]')
  saas_mode_enabled = GDK::Output.prompt('Do you want to enable Duo Features in SaaS (GitLab.com) Mode? [y/N]')
  enable_hot_reload = GDK::Output.prompt('Do you want to enable hot reload?[y/N]')

  GDK::Output.puts 'Enabling GitLab AI Gateway in GDK config...'
  GDK.config.bury!('gitlab_ai_gateway.enabled', true)
  GDK.config.save_yaml!

  GDK::Output.puts 'Updating GDK...'
  success = GDK::Command::Update.new.run

  unless success
    GDK::Output.error("Updating GDK failed. Make sure `gdk update` command succeeds in your terminal.")
    next
  end

  unless File.file?(AI_GATEWAY_ENV_FILE)
    GDK::Output.error("AI Gateway env file was not found at #{AI_GATEWAY_ENV_FILE}. Retrying `rake setup_ai_gateway` command might resolve the issue.")
    next
  end

  GDK::Output.puts 'Setting up Anthropic API key...'
  update_env_file(AI_GATEWAY_ENV_FILE, 'ANTHROPIC_API_KEY', anthropic_key)
  update_env_file(AI_GATEWAY_ENV_FILE, 'AIGW_AUTH__BYPASS_EXTERNAL', 'true')

  GDK::Output.puts 'Setting up Google Cloud...'
  GDK.make('gitlab-ai-gateway-gcloud-setup')

  if enable_debug.match?(/\Ay(?:es)*\z/i)
    DEBUG_VARS.each do |key, value|
      update_env_file(AI_GATEWAY_ENV_FILE, key, value)
    end

    GDK::Output.puts 'Debug variables have been set.'
  end

  if enable_hot_reload.match?(/\Ay(?:es)*\z/i)
    update_env_file(AI_GATEWAY_ENV_FILE, 'AIGW_FASTAPI__RELOAD', 'true')
    GDK::Output.puts 'Hot reload has been enabled.'
  end

  env_runit_contents = File.exist?(ENV_RUNIT_FILE) ? File.read(ENV_RUNIT_FILE) : ''

  unless env_runit_contents.include?('export GITLAB_SIMULATE_SAAS=') && env_runit_contents.include?('export AI_GATEWAY_URL=http://0.0.0.0:5052')
    File.open(ENV_RUNIT_FILE, 'a') do |f|
      new_content = []
      new_content << "\n# Added by GitLab AI Gateway setup" unless env_runit_contents.empty?

      unless env_runit_contents.include?('export GITLAB_SIMULATE_SAAS')
        new_content << if saas_mode_enabled.match?(/\Ay(?:es)*\z/i)
                         'export GITLAB_SIMULATE_SAAS=1'
                       else
                         'export GITLAB_SIMULATE_SAAS=0'
                       end
      end

      env_runit_contents.include?('export AI_GATEWAY_URL=http://127.0.0.1:5052') ||
        (new_content << 'export AI_GATEWAY_URL=http://127.0.0.1:5052')

      f.write("#{new_content.join("\n")}\n")
      GDK::Output.puts "Updated env.runit file with #{new_content}"
    end
  end

  GDK::Output.puts 'Restarting services...'
  GDK::Command::Restart.new.run

  GDK::Output.puts 'GDK AI Gateway setup complete'
  GDK::Output.puts "Access AI Gateway docs at the url listed in 'gdk status'"
  GDK::Output.puts 'For more information, see https://docs.gitlab.com/ee/development/ai_features/index.html'

  GDK::Output.success('Done')
end