has_aux_table/lib/has_aux_table.rb

# typed: true
# frozen_string_literal: true

require "sorbet-runtime"
require "active_record"
require "active_record/base"
require "active_support"
require "active_support/concern"
require "active_model/attribute_set"

require_relative "has_aux_table/version"
require_relative "has_aux_table/key_type"
require_relative "has_aux_table/util"
require_relative "has_aux_table/relation_extensions"
require_relative "has_aux_table/aux_table_config"
require_relative "has_aux_table/migration_extensions"

module HasAuxTable
  extend T::Sig
  extend T::Helpers
  extend ActiveSupport::Concern

  included do
    T.bind(self, T.class_of(ActiveRecord::Base))
    before_create do
      T.bind(self, ActiveRecord::Base)
      T.unsafe(self).type ||= self.class.name
    end
  end

  module ClassMethods
    extend T::Sig
    extend T::Helpers
    requires_ancestor { T.class_of(ActiveRecord::Base) }

    include RelationExtensions

    # Main DSL method for defining auxiliary tables
    sig { params(aux_name: T.any(String, Symbol)).returns(AuxTableConfig) }
    def aux_table(aux_name)
      @aux_table_configs ||=
        T.let({}, T.nilable(T::Hash[Symbol, AuxTableConfig]))

      aux_name = aux_name.to_sym
      if @aux_table_configs.key?(aux_name)
        Kernel.raise ArgumentError,
                     "Auxiliary '#{aux_name}' on #{self.name} (table '#{self.table_name}') already exists"
      end

      @aux_table_configs[aux_name] = config = generate_aux_config(aux_name)
      setup_attribute_types_hook!(config)
      setup_load_schema_hook!(config)
      setup_initialize_hook!(config)
      setup_changed_hook!(config)
      setup_save_hook!(config)
      setup_reload_hook!(config)
      setup_attributes_hook!(config)
      setup_relation_extensions!(config)
      setup_attribute_getter_setter_hooks!(config)
      setup_enum_hook!(config)
      setup_update_counter_hook!(config)

      config
    end

    private

    # Generate auxiliary model class dynamically
    sig do
      params(
        aux_name: Symbol,
        foreign_key: KeyType,
        primary_key: KeyType
      ).returns(AuxTableConfig)
    end
    def generate_aux_config(
      aux_name,
      # The column on the aux table that points to the main table
      foreign_key: :base_table_id,
      primary_key: self.primary_key
    )
      base_table = self.table_name
      aux_table_name = :"#{base_table}_#{aux_name}_aux"

      # Generate class name (e.g., :car_aux => "CarAux")
      aux_class_name = aux_table_name.to_s.camelize
      aux_association_name = aux_table_name.to_s.singularize.to_sym

      # Ensure the class name doesn't conflict with existing constants
      if Object.const_defined?(aux_class_name)
        Object.send(:remove_const, aux_class_name)
      end

      # Get the current class for the association
      main_class = T.cast(self, T.class_of(ActiveRecord::Base))
      main_association_name = foreign_key.to_s.delete_suffix("_id").to_sym

      # Create the auxiliary model class
      aux_class =
        Class.new(ActiveRecord::Base) do
          self.table_name = aux_table_name.to_s
          self.primary_key = foreign_key

          # Define the association back to the specific STI subclass
          # Foreign key points to base STI table (e.g., vehicle_id)
          # But association is to the specific subclass (e.g., Car)
          self.belongs_to(
            main_association_name,
            class_name: main_class.name,
            foreign_key:,
            primary_key:,
            inverse_of: aux_association_name
          )
        end

      # set up has_one association to the auxiliary table
      self.has_one(
        aux_association_name,
        class_name: aux_class_name,
        foreign_key:,
        primary_key:,
        inverse_of: main_association_name,
        dependent: :destroy
      )

      # so the aux table is joined against the main table
      self.default_scope { eager_load(aux_association_name) }

      # Set the constant to make the class accessible
      Object.const_set(aux_class_name, aux_class)

      AuxTableConfig.new(
        aux_table_name:,
        aux_class:,
        main_class:,
        aux_association_name:,
        main_association_name:,
        foreign_key:,
        primary_key:
      )
    end

    sig { params(config: AuxTableConfig).void }
    def setup_attribute_types_hook!(config)
      original_method = config.main_class.method(:attribute_types)
      config
        .main_class
        .define_singleton_method(:attribute_types) do
          @aux_config_attribute_types_cache ||=
            T.let(
              {},
              T.nilable(
                T::Hash[Symbol, T::Hash[String, ActiveModel::Type::Value]]
              )
            )

          @aux_config_attribute_types_cache[config.aux_table_name] ||= begin
            original_types =
              T.let(
                original_method.call,
                T::Hash[String, ActiveModel::Type::Value]
              )

            # move 'created_at', 'updated_at' etc to the end of the list
            timestamp_types = {}
            original_types.reject! do |k, v|
              if k.end_with?("_at") && v.type == :datetime
                timestamp_types[k] = v
                original_types.delete(k)
              end
            end

            original_types.merge!(config.aux.attribute_types)
            original_types.merge!(timestamp_types)
            original_types
          end
        end

      config.main_class.attributes_for_inspect =
        Util.attributes_for_inspect(config)
    end

    # Hook into schema loading to generate attribute accessors when schema is loaded
    sig { params(config: AuxTableConfig).void }
    def setup_load_schema_hook!(config)
      # Override load_schema to also generate auxiliary attribute accessors when schema is loaded
      load_schema_method = self.method(:load_schema!)
      self.define_singleton_method(:load_schema!) do
        T.bind(
          self,
          T.all(T.class_of(ActiveRecord::Base), HasAuxTable::ClassMethods)
        )

        aux_config_load_schema!(load_schema_method, config)
      end

      self.load_schema! if self.schema_loaded?
    end

    sig { params(load_schema_method: Method, config: AuxTableConfig).void }
    def aux_config_load_schema!(load_schema_method, config)
      # first, load the main and aux table schemas like normal
      config.load_aux_schema
      result = load_schema_method.call
      aux_table_name = config.aux_table_name

      check_for_overlapping_columns!(
        aux_table_name,
        config.main.column_names,
        config.aux.column_names
      )

      # set attributes that exist on the aux table to also exist on this table
      config.aux.default_attributes.each do |name, attr|
        @default_attributes[name] = attr
      end

      # Generate attribute accessors for each auxiliary column
      config.aux.columns_hash.each do |column_name, column|
        column_name = column_name.to_sym
        if self.method_defined?(column_name.to_sym)
          raise "invariant: method #{column_name} already defined"
        end

        config.define_aux_attribute_delegate(column_name)
        config.define_aux_attribute_delegate(:"#{column_name}?")
        config.define_aux_attribute_delegate(:"#{column_name}=")
      end

      result
    end

    sig { params(config: AuxTableConfig).void }
    def setup_attribute_getter_setter_hooks!(config)
      %i[
        _read_attribute
        read_attribute
        _write_attribute
        write_attribute
      ].each do |method_name|
        method = self.instance_method(method_name)
        self.define_method(method_name) do |name, *args, **kwargs, &block|
          T.bind(self, ActiveRecord::Base)
          if config.aux.column_names.include?(name.to_s)
            target = config.aux_model_for(self)
            T.unsafe(target).send(method_name, name, *args, **kwargs, &block)
          else
            T.unsafe(method).bind(self).call(name, *args, **kwargs, &block)
          end
        end
      end
    end

    sig { params(config: AuxTableConfig).void }
    def setup_initialize_hook!(config)
      initialize_method = self.instance_method(:initialize)
      self.define_method(:initialize) do |*args, **kwargs, &block|
        T.bind(self, ActiveRecord::Base)
        if args && args.size == 1 && (arg = args.first).is_a?(Hash)
          main_args, aux_args = config.aux.partition_by_columns(args.first)
          initialize_method.bind(self).call(main_args, **kwargs, &block)
          config.aux_model_for(self).assign_attributes(aux_args)
        else
          initialize_method.bind(self).call(*args, **kwargs, &block)
          config.aux_model_for(self)
        end
      end
    end

    sig { params(config: AuxTableConfig).void }
    def setup_changed_hook!(config)
      changed_method = self.instance_method(:changed?)
      self.define_method(:changed?) do
        T.bind(self, ActiveRecord::Base)
        changed_method.bind(self).call || config.aux_model_for(self).changed?
      end

      changed_attributes_method = self.instance_method(:changed_attributes)
      self.define_method(:changed_attributes) do
        T.bind(self, ActiveRecord::Base)
        changed_attributes_method
          .bind(self)
          .call
          .merge(config.aux_model_for(self).changed_attributes)
      end
    end

    sig { params(config: AuxTableConfig).void }
    def setup_save_hook!(config)
      %i[save save!].each do |method_name|
        save_method = self.instance_method(method_name)
        self.define_method(method_name) do |*args, **kwargs, &block|
          T.bind(self, ActiveRecord::Base)
          result = save_method.bind(self).call(*args, **kwargs, &block)
          result &&=
            self
              .association(config.aux_association_name)
              .target
              .send(method_name, *args, **kwargs, &block)
          result
        end
      end
    end

    sig { params(config: AuxTableConfig).void }
    def setup_reload_hook!(config)
      refresh_object =
        lambda do |target, fresh|
          target.instance_variable_set(
            :@association_cache,
            fresh.instance_variable_get(:@association_cache)
          )

          target
            .instance_variable_get(:@association_cache)
            .each_value { |association| association.owner = target }

          target.instance_variable_set(
            :@attributes,
            fresh.instance_variable_get(:@attributes)
          )
        end

      self.define_method(:reload) do |*args|
        T.bind(self, ActiveRecord::Base)
        aux_model = config.aux_model_for(self)
        fresh_object = self.class.find(id)
        refresh_object.call(self, fresh_object)
        refresh_object.call(
          aux_model,
          fresh_object.association(config.aux_association_name).target
        )
        self
      end
    end

    sig { params(config: AuxTableConfig).void }
    def setup_attributes_hook!(config)
      attributes_method = self.instance_method(:attributes)
      self.define_method(:attributes) do |*args|
        T.bind(self, ActiveRecord::Base)
        ret = attributes_method.bind(self).call(*args)
        target = config.aux_model_for(self)
        ret.merge!(config.aux.attributes_on(target))
        ret
      end
    end

    sig do
      params(
        aux_table_name: Symbol,
        main_columns: T::Array[String],
        aux_columns: T::Array[String]
      ).void
    end
    def check_for_overlapping_columns!(
      aux_table_name,
      main_columns,
      aux_columns
    )
      # Find overlapping columns (excluding system columns and foreign keys)
      overlapping_columns =
        aux_columns.select { |col| main_columns.include?(col) }

      if overlapping_columns.any?
        column_list = overlapping_columns.map { |col| "'#{col}'" }.join(", ")
        Kernel.raise ArgumentError,
                     "Auxiliary table '#{aux_table_name}' defines column(s) #{column_list} " \
                       "that already exist(s) in main table '#{self.table_name}'. " \
                       "Auxiliary table columns must not overlap with main table columns."
      end
    end

    sig { params(config: AuxTableConfig).void }
    def setup_enum_hook!(config)
      self.define_singleton_method(
        :enum
      ) do |name = nil, values = nil, **options, &block|
        name = name.to_s
        config.aux.klass.enum(name, values, **options, &block)
        self.defined_enums.merge!(config.aux.klass.defined_enums)

        # define methods on the main class that delegate to the aux class
        config
          .aux
          .klass
          .send(:_enum_methods_module)
          .instance_methods
          .each do |method_name|
            self
              ._enum_methods_module
              .define_method(method_name) do |*args, **kwargs, &block|
                raise "not implemented"
              end
          end

        define_singleton_method(name.pluralize) do
          config.aux.klass.send(name.pluralize)
        end

        values.keys.each do |value_name|
          value_name = [options[:prefix], value_name, options[:suffix]].reject(
            &:blank?
          ).join("_")
          [
            value_name,
            "#{value_name}?",
            "#{value_name}!",
            "#{value_name}="
          ].each do |method_name|
            config.define_aux_attribute_delegate(method_name.to_sym)
          end
        end
      end
    end

    sig { params(config: AuxTableConfig).void }
    def setup_update_counter_hook!(config)
      self.define_singleton_method(:update_counters) do |id, counters|
        T.bind(self, T.class_of(ActiveRecord::Base))
        main_counters = {}
        aux_counters = {}
        opts = {}
        counters.each do |k, v|
          is_aux = config.aux.is_column?(k)
          is_main = config.main.is_column?(k)
          if !is_aux && !is_main
            opts[k] = v
          elsif is_aux
            aux_counters[k] = v
          elsif is_main
            main_counters[k] = v
          end
        end
        super(id, main_counters.merge(opts)) if main_counters.any?
        if aux_counters.any?
          config.aux.klass.update_counters(id, aux_counters.merge(opts))
        end
      end
    end
  end
  mixes_in_class_methods(ClassMethods)
end