module ActiveRecord
  module AssociationPreload
    module ClassMethods
      def find_associated_records(ids, reflection, preload_options)
        options = reflection.options
        table_name = reflection.klass.quoted_table_name

        if interface = reflection.options[:as]
          conditions = "#{reflection.klass.quoted_table_name}.#{connection.quote_column_name "#{interface}_id"} #{in_or_equals_for_ids(ids)} and #{reflection.klass.quoted_table_name}.#{connection.quote_column_name "#{interface}_type"} IN ('#{self.name}','#{self.base_class.sti_name}')"
        else
          foreign_key = reflection.primary_key_name
          conditions = "#{reflection.klass.quoted_table_name}.#{foreign_key} #{in_or_equals_for_ids(ids)}"
        end

        conditions << append_conditions(reflection, preload_options)

        reflection.klass.with_exclusive_scope do
          reflection.klass.find(:all,
                              :select => (preload_options[:select] || options[:select] || "#{table_name}.*"),
                              :include => preload_options[:include] || options[:include],
                              :conditions => [conditions, ids],
                              :joins => options[:joins],
                              :group => preload_options[:group] || options[:group],
                              :order => preload_options[:order] || options[:order])
        end
      end
      
      def preload_through_records(records, reflection, through_association)
        through_reflection = reflections[through_association]
        through_primary_key = through_reflection.primary_key_name

        if reflection.options[:source_type]
          interface = reflection.source_reflection.options[:foreign_type]
          preload_options = {:conditions => ["#{connection.quote_column_name interface} IN (?)", reflection.options[:source_type]]}
          
          records.compact!
          records.first.class.preload_associations(records, through_association, preload_options)

          # Dont cache the association - we would only be caching a subset
          through_records = []
          records.each do |record|
            proxy = record.send(through_association)

            if proxy.respond_to?(:target)
              through_records << proxy.target
              proxy.reset
            else # this is a has_one :through reflection
              through_records << proxy if proxy
            end
          end
          through_records.flatten!
        else
          records.first.class.preload_associations(records, through_association)
          through_records = records.map {|record| record.send(through_association)}.flatten
        end
        through_records.compact!
        through_records
      end
    end
  end
  
  module Associations
    class HasManyThroughAssociation < HasManyAssociation
      protected
      # added support for STI with polymorphism
      def construct_conditions
        table_name = @reflection.through_reflection.quoted_table_name
        conditions = construct_quoted_owner_attributes(@reflection.through_reflection).map do |attr, value|
          if attr =~ /_type$/
            construct_polymorphic_sql(table_name, attr)
          else
            "#{table_name}.#{attr} = #{value}"
          end
        end
        conditions << sql_conditions if sql_conditions
        
        "(" + conditions.join(') AND (') + ")"
      end
      
      # Construct attributes for associate pointing to owner.
      def construct_owner_attributes(reflection)
        if as = reflection.options[:as]
          { "#{as}_id" => @owner.id,
            "#{as}_type" => @owner.class.name.to_s }
        else
          { reflection.primary_key_name => @owner.id }
        end
      end

      # Construct attributes for :through pointing to owner and associate.
      def construct_join_attributes(associate)
        # TODO: revist this to allow it for deletion, supposing dependent option is supported
        raise ActiveRecord::HasManyThroughCantAssociateThroughHasOneOrManyReflection.new(@owner, @reflection) if [:has_one, :has_many].include?(@reflection.source_reflection.macro)
        join_attributes = construct_owner_attributes(@reflection.through_reflection).merge(@reflection.source_reflection.primary_key_name => associate.id)

        if @reflection.options[:source_type]
          join_attributes.merge!(@reflection.source_reflection.options[:foreign_type] => associate.class.name.to_s)
        end
        join_attributes
      end

      # Associate attributes pointing to owner, quoted.
      def construct_quoted_owner_attributes(reflection)
        
        if as = reflection.options[:as]
          types = [@owner.class.name.to_s, @owner.class.base_class.name.to_s].map do |clazz|
            reflection.klass.quote_value(clazz, reflection.klass.columns_hash["#{as}_type"])
          end
          { "#{as}_id" => owner_quoted_id, "#{as}_type" => types }
        elsif reflection.macro == :belongs_to
          { reflection.klass.primary_key => @owner[reflection.primary_key_name] }
        else
          { reflection.primary_key_name => owner_quoted_id }
        end
      end
      
      def construct_joins(custom_joins = nil)
        polymorphic_join = nil
        if @reflection.source_reflection.macro == :belongs_to
          reflection_primary_key = @reflection.klass.primary_key
          source_primary_key     = @reflection.source_reflection.primary_key_name
          if @reflection.options[:source_type]
            polymorphic_join = construct_polymorphic_sql(
              @reflection.through_reflection.quoted_table_name,
              "#{@reflection.source_reflection.options[:foreign_type]}",
              @reflection.options[:source_type]
            )
            polymorphic_join = "AND (#{polymorphic_join})"
          end
        else
          reflection_primary_key = @reflection.source_reflection.primary_key_name
          source_primary_key     = @reflection.through_reflection.klass.primary_key
          if @reflection.source_reflection.options[:as]
            polymorphic_join = construct_polymorphic_sql(
              @reflection.quoted_table_name,
              "#{@reflection.source_reflection.options[:as]}_type",
              @reflection.through_reflection.klass
            )
            polymorphic_join = "AND (#{polymorphic_join})"
          end
        end
        
        result = "INNER JOIN %s ON %s.%s = %s.%s %s #{@reflection.options[:joins]} #{custom_joins}" % [
          @reflection.through_reflection.quoted_table_name,
          @reflection.quoted_table_name, reflection_primary_key,
          @reflection.through_reflection.quoted_table_name, source_primary_key,
          polymorphic_join
        ]
        
        result
      end
    end
    
    class HasManyAssociation < AssociationCollection
      protected
      def construct_polymorphic_sql(table_name, attribute, clazz = @owner.class)
        clazz = clazz.to_s.constantize unless clazz.is_a?(Class)
        ancestor_classes = (clazz.ancestors.reverse - clazz.included_modules).uniq + clazz.send(:subclasses)
        in_tree = []
        while ancestor = ancestor_classes.pop
          break if ancestor == clazz.base_class.superclass
          in_tree << clazz.quote_value(ancestor.name)          
        end
        "#{table_name}.#{attribute} IN (#{in_tree.join(',')})"
      end
      
      # added support for STI with polymorphism
      def construct_sql
        case
          when @reflection.options[:finder_sql]
            @finder_sql = interpolate_sql(@reflection.options[:finder_sql])

          when @reflection.options[:as]
            @finder_sql = 
              "(#{@reflection.quoted_table_name}.#{@reflection.options[:as]}_id = #{owner_quoted_id})"
            polymorphic_conditions = construct_polymorphic_sql(@reflection.quoted_table_name, "#{@reflection.options[:as]}_type")
            @finder_sql << " AND (#{polymorphic_conditions})"
            @finder_sql << " AND (#{conditions})" if conditions
          else
            @finder_sql = "#{@reflection.quoted_table_name}.#{@reflection.primary_key_name} = #{owner_quoted_id}"
            @finder_sql << " AND (#{conditions})" if conditions
        end
        
        if @reflection.options[:counter_sql]
          @counter_sql = interpolate_sql(@reflection.options[:counter_sql])
        elsif @reflection.options[:finder_sql]
          # replace the SELECT clause with COUNT(*), preserving any hints within /* ... */
          @reflection.options[:counter_sql] = @reflection.options[:finder_sql].sub(/SELECT (\/\*.*?\*\/ )?(.*)\bFROM\b/im) { "SELECT #{$1}COUNT(*) FROM" }
          @counter_sql = interpolate_sql(@reflection.options[:counter_sql])
        else
          @counter_sql = @finder_sql
        end
      end
    end
  end
end