<commit_before>package com.sksamuel.elastic4s

import org.elasticsearch.action.search._
import org.elasticsearch.index.query.QueryBuilder
import org.elasticsearch.search.sort.SortBuilder
import org.elasticsearch.search.rescore.RescoreBuilder

/** @author <NAME> */
trait SearchDsl extends QueryDsl with FilterDsl with FacetDsl with HighlightDsl with SortDsl with SuggestionDsl {

  def find = new SearchExpectsIndex
  def select = new SearchExpectsIndex
  def search = new SearchExpectsIndex
  def search(indexes: String*): SearchDefinition = new SearchDefinition(indexes)
  class SearchExpectsIndex {
    def in(indexes: String*): SearchDefinition = new SearchDefinition(indexes)
    def in(tuple: (String, String)): SearchDefinition = new SearchDefinition(Seq(tuple._1)).types(tuple._2)
  }

  class MultiSearchDefinition(searches: Iterable[SearchDefinition])
    extends RequestDefinition(MultiSearchAction.INSTANCE) {
    def build = {
      val builder = new MultiSearchRequestBuilder(null)
      searches foreach (builder add _.build)
      builder.request()
    }
  }

  def rescore(query: QueryDefinition): RescoreDefinition = {
    new RescoreDefinition(query)
  }

  class RescoreDefinition(query: QueryDefinition) {
    val builder = RescoreBuilder.queryRescorer(query.builder)
    var windowSize = 50

    def window(size: Int): RescoreDefinition = {
      this.windowSize = size
      this
    }

    def originalQueryWeight(weight: Double): RescoreDefinition = {
      builder.setQueryWeight(weight.toFloat)
      this
    }

    def rescoreQueryWeight(weight: Double): RescoreDefinition = {
      builder.setRescoreQueryWeight(weight.toFloat)
      this
    }

    def scoreMode(scoreMode: String): RescoreDefinition = {
      builder.setScoreMode(scoreMode)
      this
    }
  }

  class SearchDefinition(indexes: Seq[String]) extends RequestDefinition(SearchAction.INSTANCE) {

    // TODO Discuss: temporarily open to pass tests
    private[elastic4s] val _builder = new SearchRequestBuilder(null).setIndices(indexes: _*)
    def build = _builder.request()

    /**
     * Adds a single string query to this search
     *
     * @param string the query string
     *
     * @return this
     */
    def query(string: String): SearchDefinition = query(new StringQueryDefinition(string))
    def query(block: => QueryDefinition): SearchDefinition = query2(block.builder)
    def query2(block: => QueryBuilder): SearchDefinition = {
      _builder.setQuery(block)
      this
    }
    def bool(block: => BoolQueryDefinition): SearchDefinition = {
      _builder.setQuery(block.builder)
      this
    }

    def filter(block: => FilterDefinition): SearchDefinition = {
      _builder.setFilter(block.builder)
      this
    }

    def facets(iterable: Iterable[FacetDefinition]): SearchDefinition = {
      iterable.foreach(facet => _builder.addFacet(facet.builder))
      this
    }
    def facets(f: FacetDefinition*): SearchDefinition = facets(f.toIterable)

    def sort(sorts: SortDefinition*): SearchDefinition = sort2(sorts.map(_.builder): _*)
    def sort2(sorts: SortBuilder*): SearchDefinition = {
      sorts.foreach(_builder addSort)
      this
    }

    def suggestions(suggestions: SuggestionDefinition*): SearchDefinition = {
      suggestions.foreach(_builder addSuggestion _.builder)
      this
    }

    /**
     * Adds a single prefix query to this search
     *
     * @param tuple - the field and prefix value
     *
     * @return this
     */
    def prefix(tuple: (String, Any)) = {
      val q = new PrefixQueryDefinition(tuple._1, tuple._2)
      _builder.setQuery(q.builder.buildAsBytes)
      this
    }

    /**
     * Adds a single regex query to this search
     *
     * @param tuple - the field and regex value
     *
     * @return this
     */
    def regex(tuple: (String, Any)) = {
      val q = new RegexQueryDefinition(tuple._1, tuple._2)
      _builder.setQuery(q.builder.buildAsBytes)
      this
    }

    def term(tuple: (String, Any)) = {
      val q = new TermQueryDefinition(tuple._1, tuple._2)
      _builder.setQuery(q.builder.buildAsBytes)
      this
    }

    def range(field: String) = {
      val q = new RangeQueryDefinition(field)
      _builder.setQuery(q.builder.buildAsBytes)
      this
    }

    def highlighting(options: HighlightOptionsDefinition, highlights: HighlightDefinition*) = {
      options._encoder.foreach(encoder => _builder.setHighlighterEncoder(encoder.elastic))
      options._tagSchema.foreach(arg => _builder.setHighlighterTagsSchema(arg.elastic))
      options._order.foreach(arg => _builder.setHighlighterOrder(arg.elastic))
      _builder.setHighlighterPostTags(options._postTags: _*)
      _builder.setHighlighterPreTags(options._preTags: _*)
      _builder.setHighlighterRequireFieldMatch(options._requireFieldMatch)
      highlights.foreach(highlight => _builder.addHighlightedField(highlight.builder))
      this
    }

    def highlighting(highlights: HighlightDefinition*): SearchDefinition = {<commit_msg>Correctly construct the SearchDefinition when calling 'highlighting(highlights: HighlightDefinition*)' API<commit_after>package com.sksamuel.elastic4s

import org.elasticsearch.action.search._
import org.elasticsearch.index.query.QueryBuilder
import org.elasticsearch.search.sort.SortBuilder
import org.elasticsearch.search.rescore.RescoreBuilder

/** @author <NAME> */
trait SearchDsl extends QueryDsl with FilterDsl with FacetDsl with HighlightDsl with SortDsl with SuggestionDsl {

  def find = new SearchExpectsIndex
  def select = new SearchExpectsIndex
  def search = new SearchExpectsIndex
  def search(indexes: String*): SearchDefinition = new SearchDefinition(indexes)
  class SearchExpectsIndex {
    def in(indexes: String*): SearchDefinition = new SearchDefinition(indexes)
    def in(tuple: (String, String)): SearchDefinition = new SearchDefinition(Seq(tuple._1)).types(tuple._2)
  }

  class MultiSearchDefinition(searches: Iterable[SearchDefinition])
    extends RequestDefinition(MultiSearchAction.INSTANCE) {
    def build = {
      val builder = new MultiSearchRequestBuilder(null)
      searches foreach (builder add _.build)
      builder.request()
    }
  }

  def rescore(query: QueryDefinition): RescoreDefinition = {
    new RescoreDefinition(query)
  }

  class RescoreDefinition(query: QueryDefinition) {
    val builder = RescoreBuilder.queryRescorer(query.builder)
    var windowSize = 50

    def window(size: Int): RescoreDefinition = {
      this.windowSize = size
      this
    }

    def originalQueryWeight(weight: Double): RescoreDefinition = {
      builder.setQueryWeight(weight.toFloat)
      this
    }

    def rescoreQueryWeight(weight: Double): RescoreDefinition = {
      builder.setRescoreQueryWeight(weight.toFloat)
      this
    }

    def scoreMode(scoreMode: String): RescoreDefinition = {
      builder.setScoreMode(scoreMode)
      this
    }
  }

  class SearchDefinition(indexes: Seq[String]) extends RequestDefinition(SearchAction.INSTANCE) {

    // TODO Discuss: temporarily open to pass tests
    private[elastic4s] val _builder = new SearchRequestBuilder(null).setIndices(indexes: _*)
    def build = _builder.request()

    /**
     * Adds a single string query to this search
     *
     * @param string the query string
     *
     * @return this
     */
    def query(string: String): SearchDefinition = query(new StringQueryDefinition(string))
    def query(block: => QueryDefinition): SearchDefinition = query2(block.builder)
    def query2(block: => QueryBuilder): SearchDefinition = {
      _builder.setQuery(block)
      this
    }
    def bool(block: => BoolQueryDefinition): SearchDefinition = {
      _builder.setQuery(block.builder)
      this
    }

    def filter(block: => FilterDefinition): SearchDefinition = {
      _builder.setFilter(block.builder)
      this
    }

    def facets(iterable: Iterable[FacetDefinition]): SearchDefinition = {
      iterable.foreach(facet => _builder.addFacet(facet.builder))
      this
    }
    def facets(f: FacetDefinition*): SearchDefinition = facets(f.toIterable)

    def sort(sorts: SortDefinition*): SearchDefinition = sort2(sorts.map(_.builder): _*)
    def sort2(sorts: SortBuilder*): SearchDefinition = {
      sorts.foreach(_builder addSort)
      this
    }

    def suggestions(suggestions: SuggestionDefinition*): SearchDefinition = {
      suggestions.foreach(_builder addSuggestion _.builder)
      this
    }

    /**
     * Adds a single prefix query to this search
     *
     * @param tuple - the field and prefix value
     *
     * @return this
     */
    def prefix(tuple: (String, Any)) = {
      val q = new PrefixQueryDefinition(tuple._1, tuple._2)
      _builder.setQuery(q.builder.buildAsBytes)
      this
    }

    /**
     * Adds a single regex query to this search
     *
     * @param tuple - the field and regex value
     *
     * @return this
     */
    def regex(tuple: (String, Any)) = {
      val q = new RegexQueryDefinition(tuple._1, tuple._2)
      _builder.setQuery(q.builder.buildAsBytes)
      this
    }

    def term(tuple: (String, Any)) = {
      val q = new TermQueryDefinition(tuple._1, tuple._2)
      _builder.setQuery(q.builder.buildAsBytes)
      this
    }

    def range(field: String) = {
      val q = new RangeQueryDefinition(field)
      _builder.setQuery(q.builder.buildAsBytes)
      this
    }

    def highlighting(options: HighlightOptionsDefinition, highlights: HighlightDefinition*) = {
      options._encoder.foreach(encoder => _builder.setHighlighterEncoder(encoder.elastic))
      options._tagSchema.foreach(arg => _builder.setHighlighterTagsSchema(arg.elastic))
      options._order.foreach(arg => _builder.setHighlighterOrder(arg.elastic))
      _builder.setHighlighterPostTags(options._postTags: _*)
      _builder.setHighlighterPreTags(options._preTags: _*)
      _builder.setHighlighterRequireFieldMatch(options._requireFieldMatch)
      highlights.foreach(highlight => _builder.addHighlightedField(highlight.builder))
      this
    }

    def highlighting(highlights: HighlightDefinition*): SearchDefinition = {
      highlights.foreach(highlight => _builder.addHighlightedField(highlight.builder))
{'fraction_non_alphanumeric': 0.09576553616937855, 'fraction_numerical': 0.003136639874534405, 'mean_word_length': 3.504635761589404, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3310267', 'n_tokens_mistral': 2934, 'n_tokens_neox': 2877, 'n_words': 628}