module Neighbor
  module Utils
    def self.validate_dimensions(value, type, expected, adapter)
      dimensions = type == :sparsevec ? value.dimensions : value.size
      dimensions *= 8 if type == :bit && [:sqlite, :mysql].include?(adapter)

      if expected && dimensions != expected
        "Expected #{expected} dimensions, not #{dimensions}"
      end
    end

    def self.validate_finite(value, type)
      case type
      when :bit, :integer
        true
      when :sparsevec
        value.values.all?(&:finite?)
      else
        value.all?(&:finite?)
      end
    end

    def self.validate(value, dimensions:, type:, adapter:)
      if (message = validate_dimensions(value, type, dimensions, adapter))
        raise Error, message
      end

      if !validate_finite(value, type)
        raise Error, "Values must be finite"
      end
    end

    def self.normalize(value, column_info:)
      return nil if value.nil?

      raise Error, "Normalize not supported for type" unless [:cube, :vector, :halfvec].include?(column_info&.type)

      norm = Math.sqrt(value.sum { |v| v * v })

      # store zero vector as all zeros
      # since NaN makes the distance always 0
      # could also throw error
      norm > 0 ? value.map { |v| v / norm } : value
    end

    def self.array?(value)
      !value.nil? && value.respond_to?(:to_a)
    end

    def self.adapter(model)
      case model.connection_db_config.adapter
      when /sqlite/i
        :sqlite
      when /mysql|trilogy/i
        model.connection_pool.with_connection { |c| c.try(:mariadb?) } ? :mariadb : :mysql
      else
        :postgresql
      end
    end

    def self.type(adapter, column_type)
      case adapter
      when :mysql
        if column_type == :binary
          :bit
        else
          column_type
        end
      else
        column_type
      end
    end

    def self.operator(adapter, column_type, distance)
      case adapter
      when :sqlite
        case distance
        when "euclidean"
          if SQLite.vec1?
            "vec1_l2_distance"
          elsif SQLite.sqlite_vec?
            "vec_distance_L2"
          else
            "neighbor_l2_distance"
          end
        when "cosine"
          if SQLite.vec1?
            "vec1_cos_distance"
          elsif SQLite.sqlite_vec?
            "vec_distance_cosine"
          else
            "neighbor_cosine_distance"
          end
        when "taxicab"
          if SQLite.sqlite_vec?
            "vec_distance_L1"
          else
            "neighbor_l1_distance"
          end
        when "inner_product"
          "neighbor_max_inner_product"
        when "hamming"
          if SQLite.sqlite_vec?
            "vec_distance_hamming"
          else
            "neighbor_hamming_distance"
          end
        when "jaccard"
          "neighbor_jaccard_distance"
        end
      when :mariadb
        case column_type
        when :vector
          case distance
          when "euclidean"
            "VEC_DISTANCE_EUCLIDEAN"
          when "cosine"
            "VEC_DISTANCE_COSINE"
          end
        when :integer
          case distance
          when "hamming"
            "BIT_COUNT"
          end
        else
          raise ArgumentError, "Unsupported type: #{column_type}"
        end
      when :mysql
        case column_type
        when :vector
          case distance
          when "cosine"
            "COSINE"
          when "euclidean"
            "EUCLIDEAN"
          end
        when :binary
          case distance
          when "hamming"
            "BIT_COUNT"
          end
        else
          raise ArgumentError, "Unsupported type: #{column_type}"
        end
      else
        case column_type
        when :bit
          case distance
          when "hamming"
            "<~>"
          when "jaccard"
            "<%>"
          when "hamming2"
            "#"
          end
        when :vector, :halfvec, :sparsevec
          case distance
          when "inner_product"
            "<#>"
          when "cosine"
            "<=>"
          when "euclidean"
            "<->"
          when "taxicab"
            "<+>"
          end
        when :cube
          case distance
          when "taxicab"
            "<#>"
          when "chebyshev"
            "<=>"
          when "euclidean", "cosine"
            "<->"
          end
        else
          raise ArgumentError, "Unsupported type: #{column_type}"
        end
      end
    end

    def self.order(adapter, type, operator, quoted_attribute, query)
      case adapter
      when :sqlite
        if operator.start_with?("neighbor")
          if type == :bit
            "#{operator}(#{quoted_attribute}, #{query})"
          else
            "#{operator}(#{quoted_attribute}, #{query}, #{type == :int8 ? 1 : 0})"
          end
        elsif type == :int8
          "#{operator}(vec_int8(#{quoted_attribute}), vec_int8(#{query}))"
        elsif type == :bit
          "#{operator}(vec_bit(#{quoted_attribute}), vec_bit(#{query}))"
        else
          "#{operator}(#{quoted_attribute}, #{query})"
        end
      when :mariadb
        if operator == "BIT_COUNT"
          "BIT_COUNT(#{quoted_attribute} ^ #{query})"
        else
          "#{operator}(#{quoted_attribute}, #{query})"
        end
      when :mysql
        if operator == "BIT_COUNT"
          "BIT_COUNT(#{quoted_attribute} ^ #{query})"
        elsif operator == "COSINE"
          "DISTANCE(#{quoted_attribute}, #{query}, 'COSINE')"
        else
          "DISTANCE(#{quoted_attribute}, #{query}, 'EUCLIDEAN')"
        end
      else
        if operator == "#"
          "bit_count(#{quoted_attribute} # #{query})"
        else
          "#{quoted_attribute} #{operator} #{query}"
        end
      end
    end

    def self.normalize_required?(adapter, column_type)
      case adapter
      when :postgresql
        column_type == :cube
      else
        false
      end
    end
  end
end
