///|
/// One document occurrence in a posting list.
pub(all) struct Posting {
  doc_id : DocId
  term_freq : Int
  positions : ReadOnlyArray[Int]
} derive(Eq, @debug.Debug)

///|
struct FieldLength {
  field_id : FieldId
  mut length : Int
}

///|
struct FieldStats {
  field_id : FieldId
  mut document_count : Int
  mut total_length : Int
}

///|
priv struct FieldPositionState {
  field_id : FieldId
  mut next_position : Int
}

///|
/// Mutable builder for one in-memory segment.
pub struct SegmentWriter[A] {
  tokenizer : A
  tokenizer_manager : TokenizerManager?
  schema : Schema?
  terms : Array[Term]
  posting_doc_ids : Array[Array[DocId]]
  posting_term_freqs : Array[Array[Int]]
  posting_positions : Array[Array[Array[Int]]]
  stored_documents : Array[StoredDocument]
  document_field_lengths : Array[Array[FieldLength]]
  field_stats : Array[FieldStats]
  mut next_doc_id : Int
}

///|
pub fn[A] SegmentWriter::new(tokenizer : A) -> SegmentWriter[A] {
  {
    tokenizer,
    tokenizer_manager: None,
    schema: None,
    terms: [],
    posting_doc_ids: [],
    posting_term_freqs: [],
    posting_positions: [],
    stored_documents: [],
    document_field_lengths: [],
    field_stats: [],
    next_doc_id: 0,
  }
}

///|
/// Creates a writer that applies the schema's indexed and stored text options.
pub fn[A] SegmentWriter::with_schema(
  schema : Schema,
  tokenizer : A,
) -> SegmentWriter[A] {
  {
    tokenizer,
    tokenizer_manager: None,
    schema: Some(schema),
    terms: [],
    posting_doc_ids: [],
    posting_term_freqs: [],
    posting_positions: [],
    stored_documents: [],
    document_field_lengths: [],
    field_stats: [],
    next_doc_id: 0,
  }
}

///|
fn Document::snapshot_with_schema(
  self : Document,
  schema : Schema,
) -> StoredDocument {
  let field_ids : Array[FieldId] = []
  let texts : Array[String] = []
  for index in 0.. field_id
      None => abort("document field identifier snapshot is inconsistent")
    }
    if schema.is_stored(field_id) {
      field_ids.push(field_id)
      match self.text_at(index) {
        Some(text) => texts.push(text)
        None => abort("document text snapshot is inconsistent")
      }
    }
  }
  StoredDocument::from_fields(field_ids, texts)
}

///|
/// Creates a schema-aware writer that resolves each indexed field's Tokenizer
/// pipeline through the supplied TokenizerManager before accepting documents.
pub fn SegmentWriter::with_schema_and_tokenizers(
  schema : Schema,
  tokenizer_manager : TokenizerManager,
) -> SegmentWriter[WhitespaceAnalyzer] raise AnalysisError {
  for field_index in 0.. name
        None => abort("schema field disappeared during validation")
      }
      if !tokenizer_manager.contains(tokenizer_name) {
        raise AnalysisError::UnknownTokenizer(tokenizer_name)
      }
    }
  }
  let tokenizer_snapshot = tokenizer_manager.snapshot()
  {
    tokenizer: WhitespaceAnalyzer::new(),
    tokenizer_manager: Some(tokenizer_snapshot),
    schema: Some(schema),
    terms: [],
    posting_doc_ids: [],
    posting_term_freqs: [],
    posting_positions: [],
    stored_documents: [],
    document_field_lengths: [],
    field_stats: [],
    next_doc_id: 0,
  }
}

///|
fn[A] SegmentWriter::add_occurrence(
  self : SegmentWriter[A],
  term : Term,
  doc_id : DocId,
  position : Int,
) -> Unit {
  match self.terms.search_by(indexed_term => indexed_term == term) {
    Some(term_index) => {
      let posting_index = self.posting_doc_ids[term_index].length() - 1
      if posting_index >= 0 &&
        self.posting_doc_ids[term_index][posting_index] == doc_id {
        self.posting_term_freqs[term_index][posting_index] += 1
        self.posting_positions[term_index][posting_index].push(position)
      } else {
        self.posting_doc_ids[term_index].push(doc_id)
        self.posting_term_freqs[term_index].push(1)
        self.posting_positions[term_index].push([position])
      }
    }
    None => {
      self.terms.push(term)
      self.posting_doc_ids.push([doc_id])
      self.posting_term_freqs.push([1])
      self.posting_positions.push([[position]])
    }
  }
}

///|
fn field_position_state(
  states : Array[FieldPositionState],
  field_id : FieldId,
) -> Int {
  match states.search_by(state => state.field_id == field_id) {
    Some(index) => index
    None => {
      states.push({ field_id, next_position: 0 })
      states.length() - 1
    }
  }
}

///|
fn record_document_field_length(
  lengths : Array[FieldLength],
  field_id : FieldId,
  token_count : Int,
) -> Unit {
  match lengths.search_by(entry => entry.field_id == field_id) {
    Some(index) => lengths[index].length += token_count
    None => lengths.push({ field_id, length: token_count })
  }
}

///|
fn[A] SegmentWriter::record_field_stats(
  self : SegmentWriter[A],
  lengths : Array[FieldLength],
) -> Unit {
  for entry in lengths {
    match
      self.field_stats.search_by(stats => stats.field_id == entry.field_id) {
      Some(index) => {
        self.field_stats[index].document_count += 1
        self.field_stats[index].total_length += entry.length
      }
      None =>
        self.field_stats.push({
          field_id: entry.field_id,
          document_count: 1,
          total_length: entry.length,
        })
    }
  }
}

///|
/// Adds a document and returns its sequential segment-local identifier.
pub fn[A : @analysis.Tokenizer] SegmentWriter::add_document(
  self : SegmentWriter[A],
  document : Document,
) -> DocId {
  let doc_id = DocId::new(self.next_doc_id)
  let stored_document = match self.schema {
    Some(schema) => document.snapshot_with_schema(schema)
    None => document.snapshot()
  }
  let field_lengths : Array[FieldLength] = []
  let field_positions : Array[FieldPositionState] = []
  for field_index in 0.. schema.is_indexed(field_id)
      None => true
    }
    if !should_index {
      continue
    }
    let token_stream = match (self.schema, self.tokenizer_manager) {
      (Some(schema), Some(manager)) => {
        let tokenizer_name = match schema.tokenizer_name(field_id) {
          Some(name) => name
          None => abort("document contains an unknown schema field")
        }
        manager.token_stream(tokenizer_name, document.texts[field_index]) catch {
          error => abort(error.to_string())
        }
      }
      _ => self.tokenizer.token_stream(document.texts[field_index])
    }
    let state_index = field_position_state(field_positions, field_id)
    let position_base = field_positions[state_index].next_position
    let mut token_count = 0
    let mut maximum_position = -1
    let mut previous_position = -1
    while token_stream.advance() {
      match token_stream.token() {
        Some(token) => {
          guard token.position >= 0 &&
            token.position >= previous_position &&
            token.position_length > 0 else {
            abort(
              "token positions must be non-decreasing and position_length must be positive",
            )
          }
          token_count += 1
          let token_final_position = token.position + token.position_length - 1
          guard token_final_position >= token.position else {
            abort("token end position overflowed")
          }
          if token_final_position > maximum_position {
            maximum_position = token_final_position
          }
          self.add_occurrence(
            Term::new(field_id, token.text),
            doc_id,
            position_base + token.position,
          )
          previous_position = token.position
        }
        None => ()
      }
    }
    record_document_field_length(field_lengths, field_id, token_count)
    if maximum_position >= 0 {
      field_positions[state_index].next_position = position_base +
        maximum_position +
        2
    }
  }
  self.stored_documents.push(stored_document)
  self.record_field_stats(field_lengths)
  self.document_field_lengths.push(field_lengths)
  self.next_doc_id += 1
  doc_id
}

///|
/// Immutable in-memory index for one batch of documents.
pub struct Segment {
  schema : Schema?
  indexed_terms : ReadOnlyArray[Term]
  posting_lists : ReadOnlyArray[ReadOnlyArray[Posting]]
  stored_documents : ReadOnlyArray[StoredDocument]
  document_field_lengths : ReadOnlyArray[ReadOnlyArray[FieldLength]]
  field_stats : ReadOnlyArray[FieldStats]
  document_count : Int
}

///|
fn freeze_positions(positions : Array[Int]) -> ReadOnlyArray[Int] {
  let frozen : Array[Int] = []
  for position in positions {
    frozen.push(position)
  }
  ReadOnlyArray::from_array(frozen)
}

///|
fn freeze_field_lengths(
  lengths : Array[FieldLength],
) -> ReadOnlyArray[FieldLength] {
  let frozen : Array[FieldLength] = []
  for entry in lengths {
    frozen.push(entry)
  }
  ReadOnlyArray::from_array(frozen)
}

///|
/// Finishes the current batch as an immutable in-memory segment.
pub fn[A] SegmentWriter::finish(self : SegmentWriter[A]) -> Segment {
  let posting_lists : Array[ReadOnlyArray[Posting]] = []
  for term_index in 0.. term)),
    posting_lists: ReadOnlyArray::from_array(posting_lists),
    stored_documents: ReadOnlyArray::from_array(self.stored_documents),
    document_field_lengths: ReadOnlyArray::from_array(document_field_lengths),
    field_stats: ReadOnlyArray::from_array(self.field_stats),
    document_count: self.next_doc_id,
  }
}

///|
/// Returns the schema used by this segment, if it was built with one.
pub fn Segment::schema(self : Segment) -> Schema? {
  self.schema
}

///|
pub fn Segment::doc_count(self : Segment) -> Int {
  self.document_count
}

///|
/// Returns all indexed terms in deterministic first-seen order.
pub fn Segment::terms(self : Segment) -> ReadOnlyArray[Term] {
  self.indexed_terms
}

///|
/// Returns the immutable stored document for a segment-local document ID.
pub fn Segment::doc(self : Segment, doc_id : DocId) -> StoredDocument? {
  if doc_id.value < 0 || doc_id.value >= self.stored_documents.length() {
    None
  } else {
    Some(self.stored_documents[doc_id.value])
  }
}

///|
/// Returns the analyzed token count for one document field.
pub fn Segment::field_length(
  self : Segment,
  doc_id : DocId,
  field_id : FieldId,
) -> Int {
  if doc_id.value < 0 || doc_id.value >= self.document_field_lengths.length() {
    return 0
  }
  match
    self.document_field_lengths[doc_id.value].search_by(entry => {
      entry.field_id == field_id
    }) {
    Some(index) => self.document_field_lengths[doc_id.value][index].length
    None => 0
  }
}

///|
pub fn Segment::has_field(
  self : Segment,
  doc_id : DocId,
  field_id : FieldId,
) -> Bool {
  if doc_id.value < 0 || doc_id.value >= self.document_field_lengths.length() {
    false
  } else {
    self.document_field_lengths[doc_id.value].search_by(entry => {
      entry.field_id == field_id
    })
    is Some(_)
  }
}

///|
/// Number of documents that contain the field, including empty field values.
pub fn Segment::field_doc_count(self : Segment, field_id : FieldId) -> Int {
  match self.field_stats.search_by(stats => stats.field_id == field_id) {
    Some(index) => self.field_stats[index].document_count
    None => 0
  }
}

///|
/// Average analyzed token count among documents that contain the field.
pub fn Segment::average_field_length(
  self : Segment,
  field_id : FieldId,
) -> Double {
  match self.field_stats.search_by(stats => stats.field_id == field_id) {
    Some(index) => {
      let stats = self.field_stats[index]
      if stats.document_count == 0 {
        0.0
      } else {
        stats.total_length.to_double() / stats.document_count.to_double()
      }
    }
    None => 0.0
  }
}

///|
/// Looks up a field-qualified term in this segment.
pub fn Segment::postings_for(
  self : Segment,
  term : Term,
) -> ReadOnlyArray[Posting] {
  match self.indexed_terms.search_by(indexed_term => indexed_term == term) {
    Some(index) => self.posting_lists[index]
    None => []
  }
}