module VCAP::BigintMigration
  class << self
    def opt_out?
      opt_out = VCAP::CloudController::Config.config&.get(:skip_bigint_id_migration)
      opt_out.nil? ? false : opt_out
    rescue VCAP::CloudController::Config::InvalidConfigPath
      false
    end

    def table_empty?(db, table)
      db[table].empty?
    end

    def change_pk_to_bigint(db, table)
      db.set_column_type(table, :id, :Bignum) if column_type(db, table, :id) != 'bigint'
    end

    def revert_pk_to_integer(db, table)
      db.set_column_type(table, :id, :integer) if column_type(db, table, :id) == 'bigint'
    end

    def add_bigint_column(db, table)
      db.add_column(table, :id_bigint, :Bignum, if_not_exists: true)
    end

    def drop_bigint_column(db, table)
      db.drop_column(table, :id_bigint, if_exists: true)
    end

    def create_trigger_function(db, table)
      drop_trigger_function(db, table)

      function = <<~FUNC
        BEGIN
          NEW.id_bigint := NEW.id;
          RETURN NEW;
        END;
      FUNC
      db.create_function(function_name(table), function, language: :plpgsql, returns: :trigger)
      db.create_trigger(table, trigger_name(table), function_name(table), each_row: true, events: :insert)
    end

    def drop_trigger_function(db, table)
      db.drop_trigger(table, trigger_name(table), if_exists: true)
      db.drop_function(function_name(table), if_exists: true)
    end

    def backfill(logger, db, table, batch_size: 10_000, iterations: -1)
      raise "table '#{table}' does not contain column 'id_bigint'" unless column_exists?(db, table, :id_bigint)

      logger.info("starting bigint backfill on table '#{table}' (batch_size: #{batch_size}, iterations: #{iterations})")
      loop do
        updated_rows = backfill_batch(db, table, :id, :id_bigint, batch_size)
        logger.info("updated #{updated_rows} rows")
        iterations -= 1 if iterations > 0
        break if updated_rows < batch_size || iterations == 0
      end
      logger.info("finished bigint backfill on table '#{table}'")
    end

    def migration_completed?(db, table)
      column_type(db, table, :id) == 'bigint'
    end

    def migration_skipped?(db, table)
      !column_exists?(db, table, :id_bigint)
    end

    def add_check_constraint(db, table)
      return if has_check_constraint?(db, table)

      constraint_name = check_constraint_name(table)
      db.alter_table(table) do
        add_constraint(name: constraint_name, not_valid: true) do
          Sequel.lit('id_bigint IS NOT NULL AND id_bigint = id')
        end
      end
    end

    def drop_check_constraint(db, table)
      return unless has_check_constraint?(db, table)

      constraint_name = check_constraint_name(table)
      db.alter_table(table) do
        drop_constraint(constraint_name)
      end
    end

    def validate_check_constraint(db, table)
      db.run("ALTER TABLE #{table} VALIDATE CONSTRAINT #{check_constraint_name(table)}")
    end

    def has_check_constraint?(db, table)
      check_constraint_exists?(db, table, check_constraint_name(table))
    end

    def drop_pk_column(db, table)
      db.drop_column(table, :id, if_exists: true)
    end

    def add_id_column(db, table)
      db.add_column(table, :id, :integer, if_not_exists: true)
    end

    def rename_bigint_column(db, table)
      db.rename_column(table, :id_bigint, :id) if column_exists?(db, table, :id_bigint) && !column_exists?(db, table, :id)
    end

    def revert_bigint_column_name(db, table)
      db.rename_column(table, :id, :id_bigint) if column_exists?(db, table, :id) && column_type(db, table, :id) == 'bigint' && !column_exists?(db, table, :id_bigint)
    end

    def add_pk_constraint(db, table)
      return if db.primary_key(table) == 'id'

      db.alter_table(table) do
        add_primary_key([:id])
      end
    end

    def drop_pk_constraint(db, table)
      return unless db.primary_key(table) == 'id'

      constraint_name = pk_constraint_name(table)
      db.alter_table(table) do
        drop_constraint(constraint_name)
        set_column_allow_null(:id, true)
      end
    end

    def add_timestamp_pk_index(db, table)
      db.add_index(table, %i[timestamp id], name: timestamp_id_index_name(table), unique: false, if_not_exists: true, concurrently: true)
    end

    def drop_timestamp_pk_index(db, table)
      db.drop_index(table, %i[timestamp id], name: timestamp_id_index_name(table), if_exists: true)
    end

    def set_pk_as_identity_with_correct_start_value(db, table)
      return if column_attribute(db, table, :id, :auto_increment) == true

      block = <<~BLOCK
        DO $$
        DECLARE
          max_id BIGINT;
        BEGIN
          SELECT COALESCE(MAX(id), 0) + 1 INTO max_id FROM #{table};

          EXECUTE format('ALTER TABLE #{table} ALTER COLUMN id ADD GENERATED BY DEFAULT AS IDENTITY (START WITH %s)', max_id);
        END $$;
      BLOCK
      db.run(block)
    end

    def drop_identity(db, table)
      db.run("ALTER TABLE #{table} ALTER COLUMN id DROP IDENTITY IF EXISTS")
    end

    def backfill_id(db, table)
      batch_size = 10_000
      loop do
        updated_rows = backfill_batch(db, table, :id_bigint, :id, batch_size)
        break if updated_rows < batch_size
      end
    end

    private

    def column_attribute(db, table, column, attribute)
      db.schema(table).find { |col, _| col == column }&.dig(1, attribute)
    end

    def column_type(db, table, column)
      column_attribute(db, table, column, :db_type)
    end

    def function_name(table)
      :"#{table}_set_id_bigint_on_insert"
    end

    def trigger_name(table)
      :"trigger_#{function_name(table)}"
    end

    def column_exists?(db, table, column)
      db[table].columns.include?(column)
    end

    def check_constraint_name(table)
      :"#{table}_check_id_bigint_matches_id"
    end

    def check_constraint_exists?(db, table, constraint_name)
      db.check_constraints(table).include?(constraint_name)
    end

    def pk_constraint_name(table)
      :"#{table}_pkey"
    end

    def timestamp_id_index_name(table)
      :"#{table}_timestamp_id_index"
    end

    def backfill_batch(db, table, from_column, to_column, batch_size)
      db.
        from(table, :batch).
        with(:batch, db[table].select(from_column).where(to_column => nil).order(from_column).limit(batch_size).for_update.skip_locked).
        where(Sequel.qualify(table, from_column) => :"batch__#{from_column}").
        update(to_column => :"batch__#{from_column}")
    end
  end
end
