///|
/// A small mutable vector collection that owns ingestion and index rebuilding.
///
/// `VectorCollection` is intended for applications that receive documents over
/// time. It keeps the public lifecycle explicit: upsert, remove, rebuild, then
/// search. The implementation uses FlatIndex as a correctness-first baseline;
/// callers can export the validated corpus to IVF, KD-Tree, or LSH when scale
/// justifies an approximate or partitioned index.
pub struct VectorCollection {
  documents : Array[Document]
  mut index : FlatIndex
}

///|
/// Create an empty collection.
pub fn VectorCollection::new() -> VectorCollection {
  { documents: [], index: FlatIndex::new() }
}

///|
/// Number of documents currently stored.
pub fn VectorCollection::length(self : VectorCollection) -> Int {
  self.documents.length()
}

///|
/// Return whether the collection has no documents.
pub fn VectorCollection::is_empty(self : VectorCollection) -> Bool {
  self.documents.length() == 0
}

///|
/// Return a shallow document snapshot in insertion order.
pub fn VectorCollection::documents(self : VectorCollection) -> Array[Document] {
  let result = []
  for doc in self.documents {
    result.push(doc)
  }
  result
}

///|
/// Locate a document by id.
pub fn VectorCollection::get(self : VectorCollection, id : String) -> Document? {
  for doc in self.documents {
    if doc.id == id {
      return Some(doc)
    }
  }
  None
}

///|
/// Add a document or replace the existing document with the same id.
pub fn VectorCollection::upsert(
  self : VectorCollection,
  doc : Document,
) -> Unit raise VectorError {
  if doc.id.trim().length() == 0 {
    raise IndexError("Document id must not be empty")
  }
  if doc.vector.length() == 0 {
    raise EmptyVector
  }
  if self.documents.length() > 0 {
    let dim = self.documents[0].vector.length()
    if doc.vector.length() != dim {
      raise DimensionMismatch(
        "Collection dimension is " +
        dim.to_string() +
        ", received " +
        doc.vector.length().to_string(),
      )
    }
  }
  let mut replaced = false
  for i = 0; i < self.documents.length(); i = i + 1 {
    if self.documents[i].id == doc.id {
      self.documents[i] = doc
      replaced = true
      break
    }
  }
  if !replaced {
    self.documents.push(doc)
  }
  self.rebuild()
}

///|
/// Remove a document. Returns whether an item was removed.
pub fn VectorCollection::remove(self : VectorCollection, id : String) -> Bool {
  for i = 0; i < self.documents.length(); i = i + 1 {
    if self.documents[i].id == id {
      ignore(self.documents.remove(i))
      self.rebuild()
      return true
    }
  }
  false
}

///|
/// Remove every document and reset the collection.
pub fn VectorCollection::clear(self : VectorCollection) -> Unit {
  self.documents.clear()
  self.index = FlatIndex::new()
}

///|
/// Replace all data after validating dimensions and ids.
pub fn VectorCollection::replace_all(
  self : VectorCollection,
  docs : Array[Document],
) -> Unit raise VectorError {
  if docs.length() > 0 {
    validate_documents(docs)
  }
  self.documents.clear()
  for doc in docs {
    self.documents.push(doc)
  }
  self.rebuild()
}

///|
/// Rebuild the exact baseline index from the current snapshot.
pub fn VectorCollection::rebuild(self : VectorCollection) -> Unit {
  self.index = FlatIndex::new()
  for doc in self.documents {
    self.index.add(doc)
  }
}

///|
/// Search the current collection using the exact baseline.
pub fn VectorCollection::search(
  self : VectorCollection,
  query : Array[Double],
  top_k : Int,
  metric : DistanceMetric,
  filters : Array[(String, String)],
) -> Array[SearchResult] raise VectorError {
  self.index.search(query, top_k, metric, filters)
}

///|
/// Search multiple queries using the exact baseline.
pub fn VectorCollection::search_batch(
  self : VectorCollection,
  queries : Array[Array[Double]],
  top_k : Int,
  metric : DistanceMetric,
  filters : Array[(String, String)],
) -> Array[Array[SearchResult]] raise VectorError {
  self.index.search_batch(queries, top_k, metric, filters)
}

///|
/// Return basic index statistics for observability.
pub fn VectorCollection::stats(self : VectorCollection) -> IndexStats {
  self.index.stats()
}

///|
/// Return documents whose metadata contains all requested pairs.
pub fn filter_documents(
  docs : Array[Document],
  filters : Array[(String, String)],
) -> Array[Document] {
  let result = []
  for doc in docs {
    if matches_filters(doc.metadata, filters) {
      result.push(doc)
    }
  }
  result
}

///|
/// Return all metadata values for a key, without duplicates.
pub fn metadata_values(docs : Array[Document], key : String) -> Array[String] {
  let result = []
  let seen = Map([])
  for doc in docs {
    for pair in doc.metadata {
      if pair.0 == key && !seen.contains(pair.1) {
        seen.set(pair.1, true)
        result.push(pair.1)
      }
    }
  }
  result
}

///|
/// Count documents in each metadata category.
pub fn metadata_counts(
  docs : Array[Document],
  key : String,
) -> Array[(String, Int)] {
  let counts = Map([])
  for value in metadata_values(docs, key) {
    counts.set(value, 0)
  }
  for doc in docs {
    for pair in doc.metadata {
      if pair.0 == key {
        let current = match counts.get(pair.1) {
          Some(value) => value
          None => 0
        }
        counts.set(pair.1, current + 1)
      }
    }
  }
  let result = []
  for entry in counts {
    result.push((entry.0, entry.1))
  }
  result
}