Pry::Commands.block_command('connect_heroku_db', 'Connect to remote Heroku DB.') do |uri, wrap_in_transaction|
  puts "Disconnecting current DB connection...".yellow
  ActiveRecord::Base.remove_connection
  adapter = uri.split("://")[0]
  host = uri.split("@")[1].split(":")[0]
  port = uri.split(":")[3].split("/")[0]
  username = uri.split("://")[1].split(":")[0]
  password = uri.split("://")[1].split(":")[1].split("@")[0]
  database = uri.split("/")[3]
  db_settings = {
    adapter: "postgresql",
    host: host,
    port: port,
    username: username,
    password: password,
    database: database
  }
  puts "Connecting to supplied host DB...".yellow
  ActiveRecord::Base.establish_connection(db_settings)
  puts "Checking if successfully connected...".yellow
  begin
    ActiveRecord::Base.connection.tables
    puts "Successfully connected!".green
    $old_pry_prompt = _pry_.prompt
    _pry_.prompt = [
      proc { |obj, nest_level, _| "pry-HEROKU(#{obj}):#{nest_level}> ".red },
      proc { |obj, nest_level, _| "pry-HEROKU(#{obj}):#{nest_level}* ".red }
    ]
    # puts wrap_in_transaction
    # if wrap_in_transaction == "true"
    #   ActiveRecord::Base.connection.execute("BEGIN")
    #   # ActiveRecord::Base.connection.execute("BEGIN TRANSACTION")
    #   # ActiveRecord::Base.connection.execute("SAVEPOINT heroku_connection_sp")
    #   $remote_connection = true
    #   _pry_.hooks.add_hook :after_session, 'tparreira_pry:reset_remote_connection' do
    #     puts "Exiting server: pry after_session hook resetting connection!".red
    #     _pry_.run_command("reset_connection") if $remote_connection == true
    #   end
    # end
  rescue ActiveRecord::ConnectionNotEstablished => e
    puts "Coudn't connect: \n#{e}\n".red
  rescue => e
    puts "Unexpected exception:\n#{e}\n".red
  end
end

Pry::Commands.block_command('reset_connection', 'Reset DB connection to development mode.') do
  if $old_pry_prompt != nil
    # ActiveRecord::Base.connection.execute("ROLLBACK")
    # NOTE: @tparreira - WHY DOES THIS SAY THERE'S NO TRANSACTION BLOCK?!
    # ActiveRecord::Base.connection.execute("ROLLBACK TO SAVEPOINT heroku_connection_sp")
    _pry_.prompt = $old_pry_prompt
    $old_pry_prompt = nil
  end
  ActiveRecord::Base.remove_connection
  ActiveRecord::Base.establish_connection
  puts "Restored original development DB!".green
end