module DbCharmer
  module ActiveRecord
    module Sharding

      def self.extended(model)
        model.cattr_accessor(:sharded_connection)
      end

      def shard_for(key, proxy_target = nil, &block)
        raise ArgumentError, "No sharded connection configured!" unless sharded_connection
        conn = sharded_connection.sharder.shard_for_key(key)
        on_db(conn, proxy_target, &block)
      end

      # The block passed to this function should return an object responding to key_field
      def create_with_allocated_shard(key_field = :id, &block)
        sharder = sharded_connection.sharder
        old_conn = self.db_charmer_allocated_shard_connection
        self.db_charmer_allocated_shard_connection ||= begin
            group = sharder.least_loaded_group
            shard_info = sharder.shard_info_by_group_id(group.id)
            {:group => group, :shard_info => shard_info, :config => sharder.shard_connection_config(shard_info, group.id)}
        end
        # Set global shard connection obj
        key_obj = on_db(self.db_charmer_allocated_shard_connection[:config], nil, &block)
        sharder.allocate_new_block_for_key_on_group(key_obj.send(key_field), self.db_charmer_allocated_shard_connection[:group]) if old_conn.nil?
        self.db_charmer_allocated_shard_connection = old_conn
        key_obj
      end

      # Run on default shard (if supported by the sharding method)
      def on_default_shard(proxy_target = nil, &block)
        raise ArgumentError, "No sharded connection configured!" unless sharded_connection

        if sharded_connection.support_default_shard?
          shard_for(:default, proxy_target, &block)
        else
          raise ArgumentError, "This model's sharding method does not support default shard"
        end
      end

      # Enumerate shards
      def on_each_shard(proxy_target = nil, &block)
        raise ArgumentError, "No sharded connection configured!" unless sharded_connection

        conns = sharded_connection.shard_connections
        raise ArgumentError, "This model's sharding method does not support shards enumeration" unless conns

        conns.each do |conn|
          on_db(conn, proxy_target, &block)
        end
      end

    end
  end
end