module Sinatra::ActiveRecordExtension

def self.registered(app)
  if ENV['DATABASE_URL'] && File.exist?("#{Dir.pwd}/config/database.yml")
    path = "#{Dir.pwd}/config/database.yml"
    url = ENV['DATABASE_URL']
    source = ERB.new(File.read(path)).result
    file_spec = YAML.respond_to?(:unsafe_load) ? YAML.unsafe_load(source) : YAML.load(source)
    file_spec ||= {}
    # ActiveRecord 6.1+ has moved the connection url resolver to another module
    if Gem.loaded_specs["activerecord"].version >= Gem::Version.create('6.1')
      url_spec = ActiveRecord::DatabaseConfigurations::ConnectionUrlResolver.new(url).to_hash
    else
      url_spec = ActiveRecord::ConnectionAdapters::ConnectionSpecification::ConnectionUrlResolver.new(url).to_hash
    end
    
    # if the configuration concerns only one database, and url_spec exist, url_spec will override the same key
    # if the configuration has multiple databases (Rails 6.0+ feature), url_spec is discarded
    # Following Active Record config convention
    # https://github.com/rails/rails/blob/main/activerecord/lib/active_record/database_configurations.rb#L169
    final_spec = file_spec.keys.map do |env|
      config = file_spec[env]
      if config.is_a?(Hash) && config.all? { |_k, v| v.is_a?(Hash) }
        [env, config]
      else
        [env, config.merge(url_spec)]
      end
    end.to_h
    app.set :database, final_spec
  elsif ENV['DATABASE_URL']
    app.set :database, ENV['DATABASE_URL']
  elsif File.exist?("#{Dir.pwd}/config/database.yml")
    app.set :database_file, "#{Dir.pwd}/config/database.yml"
  end
  unless defined?(Rake) || [:test, :production].include?(app.settings.environment)
    ActiveRecord::Base.logger = Logger.new(STDOUT)
  end
  app.helpers ActiveRecordHelper
  # re-connect if database connection dropped (Rails 3 only)
  app.before do
    if ActiveRecord::VERSION::MAJOR == 3
      ActiveRecord::Base.verify_active_connections! if ActiveRecord::Base.respond_to?(:verify_active_connections!)
    end
  end
  app.after do
    if ActiveRecord::VERSION::MAJOR < 7
      ActiveRecord::Base.clear_active_connections!
    else
      ActiveRecord::Base.connection_handler.clear_active_connections!
    end
  end
end

def database
  ActiveRecord::Base
end

def database=(spec)
  if spec.is_a?(Hash) and spec.symbolize_keys[environment.to_sym]
    ActiveRecord::Base.configurations = spec.stringify_keys
    ActiveRecord::Base.establish_connection(environment.to_sym)
  elsif spec.is_a?(Hash)     
    ActiveRecord::Base.configurations = {
      environment.to_s => spec.stringify_keys
    }
    ActiveRecord::Base.establish_connection(spec.stringify_keys)
  else
    if Gem.loaded_specs["activerecord"].version >= Gem::Version.create('6.0')
      ActiveRecord::Base.configurations ||= ActiveRecord::DatabaseConfigurations.new({}).resolve(spec)
    else
      ActiveRecord::Base.configurations ||= {}
      ActiveRecord::Base.configurations[environment.to_s] = ActiveRecord::ConnectionAdapters::ConnectionSpecification::ConnectionUrlResolver.new(spec).to_hash
    end
    ActiveRecord::Base.establish_connection(spec)
  end
end

def database_file=(path)
  path = File.join(root, path) if Pathname(path).relative? and root
  source = ERB.new(File.read(path)).result
  spec = YAML.respond_to?(:unsafe_load) ? YAML.unsafe_load(source) : YAML.load(source)
  spec ||= {}
  set :database, spec
end