///|
/// One scored document returned by a collector.
pub(all) struct SearchHit {
  address : DocAddress
  score : Double
} derive(@debug.Debug)

///|
/// Collector that retains at most the highest-scoring K documents.
pub struct TopKCollector {
  limit : Int
}

///|
pub fn TopKCollector::new(limit : Int) -> TopKCollector {
  { limit, }
}

///|
/// Tantivy-style score-ordered top documents collector.
pub struct TopDocsCollector {
  limit : Int
}

///|
pub fn TopDocsCollector::new(limit : Int) -> TopDocsCollector {
  { limit, }
}

///|
pub struct CountCollector {
  enabled : Bool
}

///|
pub fn CountCollector::new() -> CountCollector {
  { enabled: true }
}

///|
/// Runs count and TopDocs in one scorer pass.
pub struct MultiCollector {
  top_docs_limit : Int
  include_count : Bool
}

///|
pub fn MultiCollector::new(
  top_docs : TopDocsCollector,
  include_count : Bool,
) -> MultiCollector {
  { top_docs_limit: top_docs.limit, include_count }
}

///|
pub(all) struct MultiCollectorResult {
  top_docs : ReadOnlyArray[SearchHit]
  count : Int
} derive(@debug.Debug)

///|
pub(all) enum SortOrder {
  Ascending
  Descending
} derive(Eq, @debug.Debug)

///|
pub(all) enum MissingValueOrder {
  MissingFirst
  MissingLast
} derive(Eq, @debug.Debug)

///|
pub struct SortField {
  field_id : FieldId
  order : SortOrder
  missing : MissingValueOrder
}

///|
pub fn SortField::new(
  field_id : FieldId,
  order : SortOrder,
  missing : MissingValueOrder,
) -> SortField {
  { field_id, order, missing }
}

///|
pub struct Sort {
  fields : ReadOnlyArray[SortField]
}

///|
pub fn Sort::new(fields : Array[SortField]) -> Sort {
  guard fields.length() > 0 else { abort("sort requires at least one field") }
  let frozen : Array[SortField] = []
  for field in fields {
    frozen.push(field)
  }
  { fields: ReadOnlyArray::from_array(frozen) }
}

///|
fn compare_hits(left : SearchHit, right : SearchHit) -> Int {
  if left.score > right.score {
    -1
  } else if left.score < right.score {
    1
  } else if left.address.segment_ord < right.address.segment_ord {
    -1
  } else if left.address.segment_ord > right.address.segment_ord {
    1
  } else {
    left.address.doc_id.value.compare(right.address.doc_id.value)
  }
}

///|
/// Fixed-capacity max-heap whose root is the worst retained hit.
priv struct TopKHeap {
  hits : Array[SearchHit]
  limit : Int
}

///|
fn TopKHeap::new(limit : Int) -> TopKHeap {
  { hits: [], limit }
}

///|
fn TopKHeap::swap(self : TopKHeap, left : Int, right : Int) -> Unit {
  let temporary = self.hits[left]
  self.hits[left] = self.hits[right]
  self.hits[right] = temporary
}

///|
fn TopKHeap::bubble_up(self : TopKHeap, start : Int) -> Unit {
  let mut child = start
  while child > 0 {
    let parent = (child - 1) / 2
    if compare_hits(self.hits[child], self.hits[parent]) <= 0 {
      break
    }
    self.swap(child, parent)
    child = parent
  }
}

///|
fn TopKHeap::sift_down(self : TopKHeap) -> Unit {
  let mut parent = 0
  while true {
    let left = parent * 2 + 1
    if left >= self.hits.length() {
      break
    }
    let right = left + 1
    let mut worse_child = left
    if right < self.hits.length() &&
      compare_hits(self.hits[right], self.hits[left]) > 0 {
      worse_child = right
    }
    if compare_hits(self.hits[worse_child], self.hits[parent]) <= 0 {
      break
    }
    self.swap(parent, worse_child)
    parent = worse_child
  }
}

///|
fn TopKHeap::offer(self : TopKHeap, hit : SearchHit) -> Unit {
  if self.limit <= 0 {
    return
  }
  if self.hits.length() < self.limit {
    self.hits.push(hit)
    self.bubble_up(self.hits.length() - 1)
  } else if compare_hits(hit, self.hits[0]) < 0 {
    self.hits[0] = hit
    self.sift_down()
  }
}

///|
/// Returns the score a new hit must meet once the heap is full.
fn TopKHeap::competitive_score(self : TopKHeap) -> Double? {
  if self.limit > 0 && self.hits.length() >= self.limit {
    Some(self.hits[0].score)
  } else {
    None
  }
}

///|
fn TopKCollector::collect_into(
  self : TopKCollector,
  scorer : &Scorer,
  segment_ord : Int,
  deleted_docs : ReadOnlyArray[DocId],
  heap : TopKHeap,
) -> Unit {
  if self.limit <= 0 {
    return
  }
  while scorer.advance() {
    let doc_id = scorer.doc()
    if deleted_docs.search_by(deleted => deleted == doc_id) is None {
      heap.offer({
        address: DocAddress::new(segment_ord, doc_id),
        score: scorer.score(),
      })
    }
  }
}

///|
fn TopKCollector::finish(
  self : TopKCollector,
  heap : TopKHeap,
) -> ReadOnlyArray[SearchHit] {
  if self.limit <= 0 {
    return []
  }
  heap.hits.sort_by(compare_hits)
  let result : Array[SearchHit] = []
  for hit in heap.hits {
    result.push(hit)
  }
  ReadOnlyArray::from_array(result)
}