module ActiveRecord
  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]
          { "#{as}_id" => owner_quoted_id,
            "#{as}_type" => reflection.klass.quote_value(
              @owner.class.name.to_s,
              reflection.klass.columns_hash["#{as}_type"]) }
        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
        
        "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
        ]
      end
    end
    
    class HasManyAssociation < AssociationCollection
      protected
      def construct_polymorphic_sql(table_name, attribute, clazz = @owner.class)
        condition = []
        clazz = clazz.to_s.constantize unless clazz.is_a?(Class)
        ancestor_classes = (clazz.ancestors.reverse - clazz.included_modules).uniq + clazz.send(:subclasses)
        while ancestor = ancestor_classes.pop
          break if ancestor == clazz.base_class.superclass
          condition << "#{table_name}.#{attribute} = #{clazz.quote_value(ancestor.name)}"
        end
        condition.join(" OR ")
      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