module Sequel
  module Plugins
    # The association_pks plugin adds the association_pks and association_pks=
    # instance methods to the model class for each association added.  These
    # methods allow for easily returning the primary keys of the associated
    # objects, and easily modifying the associated objects to set the primary
    # keys to just the ones given:
    #
    #   Artist.one_to_many :albums
    #   artist = Artist[1]
    #   artist.album_pks # [1, 2, 3]
    #   artist.album_pks = [2, 4]
    #   artist.album_pks # [2, 4]
    #
    # Note that it uses the singular form of the association name. Also note
    # that the setter both associates to new primary keys not in the assocation
    # and disassociated from primary keys not provided to the method.
    #
    # This plugin makes modifications directly to the underlying tables,
    # it does not create or return any model objects, and therefore does
    # not call any callbacks.  If you have any association callbacks,
    # you probably should not use the setter methods.
    #
    # This plugin only works with singular primary keys, it does not work
    # with composite primary keys.
    # 
    # Usage:
    #
    #   # Make all model subclass *_to_many associations have association_pks
    #   # methods (called before loading subclasses)
    #   Sequel::Model.plugin :association_pks
    #
    #   # Make the Album *_to_many associations have association_pks
    #   # methods (called before the association methods)
    #   Album.plugin :association_pks
    module AssociationPks
      module ClassMethods
        private

        # Define a association_pks method using the block for the association reflection 
        def def_association_pks_getter(opts, &block)
          association_module_def(:"#{singularize(opts[:name])}_pks", &block)
        end

        # Define a association_pks= method using the block for the association reflection,
        # if the association is not read only.
        def def_association_pks_setter(opts, &block)
          association_module_def(:"#{singularize(opts[:name])}_pks=", &block) unless opts[:read_only]
        end

        # Add a getter that checks the join table for matching records and
        # a setter that deletes from or inserts into the join table.
        def def_many_to_many(opts)
          super
          def_association_pks_getter(opts) do
            _join_table_dataset(opts).filter(opts[:left_key]=>send(opts[:left_primary_key])).select_map(opts[:right_key])
          end
          def_association_pks_setter(opts) do |pks|
            checked_transaction do
              ds = _join_table_dataset(opts).filter(opts[:left_key]=>send(opts[:left_primary_key]))
              ds.exclude(opts[:right_key]=>pks).delete
              pks -= ds.select_map(opts[:right_key])
              pks.each{|pk| ds.insert(opts[:left_key]=>send(opts[:left_primary_key]), opts[:right_key]=>pk)}
            end
          end
        end

        # Add a getter that checks the association dataset and a setter
        # that updates the associated table.
        def def_one_to_many(opts)
          super
          return if opts[:type] == :one_to_one
          def_association_pks_getter(opts) do
            send(opts.dataset_method).select_map(opts[:primary_key])
          end
          def_association_pks_setter(opts) do |pks|
            checked_transaction do
              ds = send(opts.dataset_method)
              ds.unfiltered.filter(opts[:primary_key]=>pks).update(opts[:key]=>pk)
              ds.exclude(opts[:primary_key]=>pks).update(opts[:key]=>nil)
            end
          end
        end
      end
    end
  end
end
