# typed: strict

module JsonAttrHelpers
  extend T::Sig
  extend T::Helpers
  requires_ancestor { ActiveRecord::Migration }

  class JsonAttrMigrateContext
    extend T::Sig
    extend T::Helpers

    sig { returns(T::Array[[String, Symbol, Symbol]]) }
    attr_reader :migrate_columns

    sig { returns(T::Array[[String, String, String]]) }
    attr_reader :migrate_json_columns

    sig { params(table_name: String, migration: ActiveRecord::Migration).void }
    def initialize(table_name, migration)
      @table_name = table_name
      @migration = migration
      @migrate_columns = T.let([], T::Array[[String, Symbol, Symbol]])
      @migrate_json_columns = T.let([], T::Array[[String, String, String]])
    end

    sig do
      params(
        column_name: String,
        old_field_name: String,
        new_field_name: T.nilable(String),
      ).void
    end
    def migrate_json_column(column_name, old_field_name, new_field_name = nil)
      new_field_name ||= old_field_name
      @migrate_json_columns << [column_name, old_field_name, new_field_name]
    end

    sig do
      params(
        column_name: String,
        rails_column_type: Symbol,
        db_column_type: Symbol,
      ).void
    end
    def migrate_column(column_name, rails_column_type, db_column_type)
      @migrate_columns << [column_name, rails_column_type, db_column_type]
    end
  end

  sig do
    params(
      table_name: String,
      blk: T.proc.bind(JsonAttrMigrateContext).void,
    ).void
  end
  def migrate_json_table(table_name, &blk)
    T.bind(self, ActiveRecord::Migration)
    context = JsonAttrMigrateContext.new(table_name, self)
    context.instance_eval(&blk)

    reversible do |dir|
      dir.up do
        native_columns =
          context
            .migrate_columns
            .map do |column_name, rails_column_type, db_column_type|
            "'#{column_name}', COALESCE(to_jsonb(#{column_name}), 'null'::jsonb)"
          end

        json_columns =
          context
            .migrate_json_columns
            .map do |column_name, old_field_name, new_field_name|
            "'#{new_field_name}', COALESCE(#{column_name}->>'#{old_field_name}', 'null')::jsonb"
          end

        built_json_attrs = native_columns + json_columns

        execute <<~SQL
          UPDATE #{table_name}
          SET json_attributes = jsonb_strip_nulls(
            json_attributes ||
             jsonb_build_object(#{built_json_attrs.join(", ")})
          )
        SQL
      end

      dir.down do
        set_clauses = []

        context
          .migrate_json_columns
          .each do |column_name, old_field_name, new_field_name|
          execute <<~SQL.strip
          UPDATE #{table_name}
            SET #{column_name} = jsonb_set(#{column_name}, '{#{old_field_name}}', json_attributes->'#{new_field_name}')
            WHERE json_attributes->>'#{column_name}' IS NOT NULL
          SQL
        end

        context
          .migrate_columns
          .each do |column_name, rails_column_type, db_column_type|
          set_clauses << <<~SQL.strip
            #{column_name} = (json_attributes->>'#{column_name}')::#{db_column_type}
          SQL
        end

        execute <<~SQL
          UPDATE #{table_name}
            SET #{set_clauses.join(",\n")}
        SQL
      end
    end

    context
      .migrate_json_columns
      .map(&:first)
      .uniq
      .each do |column_name|
        rename_column table_name, column_name, "#{column_name}_backup"
      end

    context
      .migrate_columns
      .each do |column_name, rails_column_type, db_column_type|
      rename_column table_name, column_name, "#{column_name}_backup"
    end
  end
end
