///|
/// Index strategy selected for an application retrieval request.
pub(all) enum RetrievalStrategy {
  RetrievalExact
  RetrievalIvf
  RetrievalKdTree
  RetrievalLsh
}

///|
/// Request object shared by knowledge-base and service integrations.
pub(all) struct RetrievalOptions {
  query : Array[Double]
  top_k : Int
  metric : DistanceMetric
  filters : Array[(String, String)]
  expression : FilterExpression
  strategy : RetrievalStrategy
  nprobe : Int
}

///|
/// Construct a default exact retrieval request.
pub fn retrieval_options(
  query : Array[Double],
  top_k : Int,
  metric : DistanceMetric,
) -> RetrievalOptions {
  {
    query,
    top_k: if top_k < 0 {
      0
    } else {
      top_k
    },
    metric,
    filters: [],
    expression: MatchAll,
    strategy: RetrievalExact,
    nprobe: 1,
  }
}

///|
/// Return a copy of a request with key/value AND filters.
pub fn retrieval_with_filters(
  options : RetrievalOptions,
  filters : Array[(String, String)],
) -> RetrievalOptions {
  { ..options, filters, }
}

///|
/// Return a copy of a request with a composed metadata expression.
pub fn retrieval_with_expression(
  options : RetrievalOptions,
  expression : FilterExpression,
) -> RetrievalOptions {
  { ..options, expression, }
}

///|
/// Return a copy of a request using an approximate strategy.
pub fn retrieval_with_strategy(
  options : RetrievalOptions,
  strategy : RetrievalStrategy,
  nprobe : Int,
) -> RetrievalOptions {
  { ..options, strategy, nprobe: if nprobe < 1 { 1 } else { nprobe } }
}

///|
/// Result envelope useful for API responses and observability.
pub(all) struct RetrievalResponse {
  results : Array[SearchResult]
  strategy : RetrievalStrategy
  corpus_count : Int
  candidate_count : Int
  returned_count : Int
  filter_description : String
}

///|
pub impl Show for RetrievalResponse with fn output(self, logger) {
  logger.write_string(
    "RetrievalResponse{strategy: " +
    describe_strategy(self.strategy) +
    ", corpus: " +
    self.corpus_count.to_string() +
    ", candidates: " +
    self.candidate_count.to_string() +
    ", returned: " +
    self.returned_count.to_string() +
    ", filter: " +
    self.filter_description +
    "}",
  )
}

///|
/// Mutable in-memory knowledge base with exact and approximate retrieval paths.
pub struct KnowledgeBase {
  collection : VectorCollection
  mut ivf : IvfIndex?
  mut lsh : LshIndex?
  mut revision : Int
}

///|
/// Create an empty knowledge base.
pub fn KnowledgeBase::new() -> KnowledgeBase {
  { collection: VectorCollection::new(), ivf: None, lsh: None, revision: 0 }
}

///|
/// Return the number of indexed documents.
pub fn KnowledgeBase::length(self : KnowledgeBase) -> Int {
  self.collection.length()
}

///|
/// Return a monotonically increasing data revision.
pub fn KnowledgeBase::revision(self : KnowledgeBase) -> Int {
  self.revision
}

///|
/// Return a stable document snapshot for export or diagnostics.
pub fn KnowledgeBase::documents(self : KnowledgeBase) -> Array[Document] {
  self.collection.documents()
}

///|
/// Insert or replace one document and rebuild approximate indexes.
pub fn KnowledgeBase::upsert(
  self : KnowledgeBase,
  doc : Document,
) -> Unit raise VectorError {
  self.collection.upsert(doc)
  self.revision = self.revision + 1
  self.rebuild_approximate_indexes()
}

///|
/// Remove one document and report whether it existed.
pub fn KnowledgeBase::remove(self : KnowledgeBase, id : String) -> Bool {
  let removed = self.collection.remove(id)
  if removed {
    self.revision = self.revision + 1
    self.rebuild_approximate_indexes()
  }
  removed
}

///|
/// Replace the complete corpus after validation.
pub fn KnowledgeBase::replace_all(
  self : KnowledgeBase,
  docs : Array[Document],
) -> Unit raise VectorError {
  self.collection.replace_all(docs)
  self.revision = self.revision + 1
  self.rebuild_approximate_indexes()
}

///|
/// Delete all documents and reset index state.
pub fn KnowledgeBase::clear(self : KnowledgeBase) -> Unit {
  self.collection.clear()
  self.ivf = None
  self.lsh = None
  self.revision = self.revision + 1
}

///|
/// Execute one retrieval request.
pub fn KnowledgeBase::retrieve(
  self : KnowledgeBase,
  options : RetrievalOptions,
) -> RetrievalResponse raise VectorError {
  let docs = self.collection.documents()
  let corpus_count = docs.length()
  let filtered_docs = filter_documents_expression(docs, options.expression)
  let candidates = match options.strategy {
    RetrievalExact => {
      let index = build_flat_index(filtered_docs)
      index.search(
        options.query,
        options.top_k,
        options.metric,
        options.filters,
      )
    }
    RetrievalIvf => self.retrieve_ivf(options, filtered_docs)
    RetrievalKdTree => self.retrieve_kd(options, filtered_docs)
    RetrievalLsh => self.retrieve_lsh(options, filtered_docs)
  }
  {
    results: candidates,
    strategy: options.strategy,
    corpus_count,
    candidate_count: filtered_docs.length(),
    returned_count: candidates.length(),
    filter_description: describe_filter_expression(options.expression),
  }
}

///|
/// Execute a batch of requests in input order.
pub fn KnowledgeBase::retrieve_batch(
  self : KnowledgeBase,
  requests : Array[RetrievalOptions],
) -> Array[RetrievalResponse] raise VectorError {
  let responses = []
  for request in requests {
    responses.push(self.retrieve(request))
  }
  responses
}

///|
/// Run exact search over an expression-filtered subset.
pub fn KnowledgeBase::search_exact(
  self : KnowledgeBase,
  query : Array[Double],
  top_k : Int,
  metric : DistanceMetric,
  expression : FilterExpression,
) -> Array[SearchResult] raise VectorError {
  let request = retrieval_with_expression(
    retrieval_options(query, top_k, metric),
    expression,
  )
  let response = self.retrieve(request)
  response.results
}

///|
/// Return a health report for readiness probes.
pub fn KnowledgeBase::health(self : KnowledgeBase) -> HealthReport {
  inspect_corpus(self.collection.documents())
}

///|
/// Return index statistics for the exact collection.
pub fn KnowledgeBase::stats(self : KnowledgeBase) -> IndexStats {
  self.collection.stats()
}

///|
/// Return metadata values present in the knowledge base.
pub fn KnowledgeBase::metadata_values(
  self : KnowledgeBase,
  key : String,
) -> Array[String] {
  metadata_values(self.collection.documents(), key)
}

///|
fn KnowledgeBase::rebuild_approximate_indexes(self : KnowledgeBase) -> Unit {
  let docs = self.collection.documents()
  if docs.length() >= 4 {
    let ivf = IvfIndex::new(4, Cosine)
    ivf.build(docs) catch {
      _ => ()
    }
    self.ivf = Some(ivf)
    let lsh = LshIndex::new(6, docs[0].vector.length())
    for doc in docs {
      lsh.add(doc) catch {
        _ => ()
      }
    }
    self.lsh = Some(lsh)
  } else {
    self.ivf = None
    self.lsh = None
  }
}

///|
fn KnowledgeBase::retrieve_ivf(
  self : KnowledgeBase,
  options : RetrievalOptions,
  filtered_docs : Array[Document],
) -> Array[SearchResult] raise VectorError {
  match self.ivf {
    Some(index) => {
      let results = index.search(
        options.query,
        options.top_k,
        options.nprobe,
        options.filters,
      )
      filter_results_expression(results, options.expression)
    }
    None => {
      let index = build_flat_index(filtered_docs)
      index.search(
        options.query,
        options.top_k,
        options.metric,
        options.filters,
      )
    }
  }
}

///|
fn KnowledgeBase::retrieve_kd(
  _self : KnowledgeBase,
  options : RetrievalOptions,
  filtered_docs : Array[Document],
) -> Array[SearchResult] raise VectorError {
  let index = build_kd_tree_index(filtered_docs, options.metric)
  filter_results_expression(
    index.search(options.query, options.top_k, options.filters),
    options.expression,
  )
}

///|
fn KnowledgeBase::retrieve_lsh(
  self : KnowledgeBase,
  options : RetrievalOptions,
  filtered_docs : Array[Document],
) -> Array[SearchResult] raise VectorError {
  match self.lsh {
    Some(index) => {
      let results = index.search(options.query, options.top_k, options.filters)
      filter_results_expression(results, options.expression)
    }
    None => {
      let index = build_flat_index(filtered_docs)
      index.search(
        options.query,
        options.top_k,
        options.metric,
        options.filters,
      )
    }
  }
}

///|
fn describe_strategy(strategy : RetrievalStrategy) -> String {
  match strategy {
    RetrievalExact => "exact"
    RetrievalIvf => "ivf"
    RetrievalKdTree => "kdtree"
    RetrievalLsh => "lsh"
  }
}