module ActiveRecord::Turntable
  module ActiveRecordExt
    module Persistence
      extend ActiveSupport::Concern

      ::ActiveRecord::Persistence.class_eval do
        # @note Override to add sharding scope on reloading
        def reload(options = nil)
          self.class.connection.clear_query_cache

          finder_scope = if turntable_enabled? && self.class.primary_key != self.class.turntable_shard_key.to_s
                           self.class.unscoped.where(self.class.turntable_shard_key => self.send(turntable_shard_key))
                         else
                           self.class.unscoped
                         end

          fresh_object =
            if options && options[:lock]
              finder_scope.lock(options[:lock]).find(id)
            else
              finder_scope.find(id)
            end

          @attributes = fresh_object.instance_variable_get("@attributes")
          @new_record = false
          self
        end

        # @note Override to add sharding scope on `touch`
        # rubocop:disable Style/UnlessElse
        def touch(*names, time: nil)
          raise ActiveRecord::ActiveRecordError, "cannot touch on a new record object" unless persisted?

          time ||= current_time_from_proper_timezone
          attributes = timestamp_attributes_for_update_in_model
          attributes.concat(names)

          unless attributes.empty?
            changes = {}

            attributes.each do |column|
              column = column.to_s
              changes[column] = write_attribute(column, time)
            end

            clear_attribute_changes(changes.keys)
            primary_key = self.class.primary_key
            scope = if turntable_enabled? && primary_key != self.class.turntable_shard_key.to_s
                      self.class.unscoped.where(self.class.turntable_shard_key => _read_attribute(turntable_shard_key))
                    else
                      self.class.unscoped
                    end
            scope = scope.where(primary_key => _read_attribute(primary_key))

            if locking_enabled?
              locking_column = self.class.locking_column
              scope = scope.where(locking_column => _read_attribute(locking_column))
              changes[locking_column] = increment_lock
            end

            result = scope.update_all(changes) == 1

            if !result && locking_enabled?
              raise ActiveRecord::StaleObjectError.new(self, "touch")
            end

            result
          else
            true
          end
        end
        # rubocop:enable Style/UnlessElse

        # @note Override to add sharding scope on `update_columns`
        def update_columns(attributes)
          raise ActiveRecord::ActiveRecordError, "cannot update a new record" if new_record?
          raise ActiveRecord::ActiveRecordError, "cannot update a destroyed record" if destroyed?

          attributes.each_key do |key|
            verify_readonly_attribute(key.to_s)
          end

          update_scope = if turntable_enabled? && self.class.primary_key != self.class.turntable_shard_key.to_s
                           self.class.unscoped.where(self.class.turntable_shard_key => self.send(turntable_shard_key))
                         else
                           self.class.unscoped
                         end

          updated_count = update_scope.where(self.class.primary_key => id).update_all(attributes)

          attributes.each do |k, v|
            raw_write_attribute(k, v)
          end

          updated_count == 1
        end

        private

          # @note Override to add sharding scope on destroying
          def relation_for_destroy
            klass = self.class
            relation = klass.unscoped.where(klass.primary_key => id)

            if klass.turntable_enabled? && klass.primary_key != klass.turntable_shard_key.to_s
              relation = relation.where(klass.turntable_shard_key => self[klass.turntable_shard_key])
            end
            relation
          end

          # @note Override to add sharding scope on updating
          def _update_record(attribute_names = self.attribute_names)
            klass = self.class
            attributes_values = arel_attributes_with_values_for_update(attribute_names)
            if attributes_values.empty?
              0
            else
              scope = if klass.turntable_enabled? && (klass.primary_key != klass.turntable_shard_key.to_s)
                        klass.unscoped.where(klass.turntable_shard_key => self.send(turntable_shard_key))
                      end
              klass.unscoped._update_record attributes_values, id, id_was, scope
            end
          end
      end
    end
  end
end