///|
/// Text chunk selected for a retrieval-augmented generation prompt.
pub(all) struct ContextChunk {
  id : String
  text : String
  score : Double
  metadata : Array[(String, String)]
}

///|
/// Prompt-ready context with source attribution.
pub(all) struct RetrievalContext {
  chunks : Array[ContextChunk]
  text : String
  source_ids : Array[String]
  total_characters : Int
}

///|
pub impl Show for RetrievalContext with fn output(self, logger) {
  logger.write_string(
    "RetrievalContext{chunks: " +
    self.chunks.length().to_string() +
    ", sources: " +
    self.source_ids.length().to_string() +
    ", characters: " +
    self.total_characters.to_string() +
    "}",
  )
}

///|
/// Read the conventional text/content/title metadata fields from a document.
pub fn document_text(doc : Document) -> String {
  let content = metadata_first(doc.metadata, "text")
  if content.length() > 0 {
    return content
  }
  let body = metadata_first(doc.metadata, "content")
  if body.length() > 0 {
    return body
  }
  let title = metadata_first(doc.metadata, "title")
  if title.length() > 0 {
    return title
  }
  let parts = []
  for pair in doc.metadata {
    parts.push(pair.0 + "=" + pair.1)
  }
  parts.join("; ")
}

///|
/// Find a document for a ranked result id.
pub fn document_for_result(
  result : SearchResult,
  docs : Array[Document],
) -> Document? {
  for doc in docs {
    if doc.id == result.id {
      return Some(doc)
    }
  }
  None
}

///|
/// Build a bounded context bundle from ranked results.
pub fn compose_retrieval_context(
  results : Array[SearchResult],
  docs : Array[Document],
  character_budget : Int,
  separator : String,
) -> RetrievalContext {
  let budget = if character_budget < 0 { 0 } else { character_budget }
  let chunks = []
  let source_ids = []
  let rendered = []
  let mut total = 0
  for result in results {
    match document_for_result(result, docs) {
      Some(doc) => {
        let text = document_text(doc)
        let block = "[" + doc.id + "]\n" + text
        let separator_cost = if rendered.length() == 0 {
          0
        } else {
          separator.length()
        }
        if total + separator_cost + block.length() <= budget ||
          (chunks.length() == 0 && total == 0) {
          chunks.push({
            id: doc.id,
            text,
            score: result.score,
            metadata: doc.metadata,
          })
          source_ids.push(doc.id)
          rendered.push(block)
          total = total + separator_cost + block.length()
        }
      }
      None => ()
    }
  }
  {
    chunks,
    text: rendered.join(separator),
    source_ids,
    total_characters: total,
  }
}

///|
/// Split a query into normalized non-empty tokens.
pub fn query_terms(query : String) -> Array[String] {
  let terms = []
  for raw in query.split(" ") {
    let term = raw.trim().to_owned()
    if term.length() > 0 {
      terms.push(term)
    }
  }
  terms
}

///|
/// Count query terms that occur in a document's textual metadata.
pub fn lexical_match_count(query : String, doc : Document) -> Int {
  let text = document_text(doc)
  let mut count = 0
  for term in query_terms(query) {
    if text.contains(term) {
      count = count + 1
    }
  }
  count
}

///|
/// Return lexical overlap in the range [0, 1].
pub fn lexical_overlap(query : String, doc : Document) -> Double {
  let terms = query_terms(query)
  if terms.length() == 0 {
    return 0.0
  }
  lexical_match_count(query, doc).to_double() / terms.length().to_double()
}

///|
/// Combine semantic and lexical relevance scores.
pub fn hybrid_score(
  semantic_score : Double,
  lexical_score : Double,
  semantic_weight : Double,
) -> Double {
  let weight = if semantic_weight < 0.0 {
    0.0
  } else if semantic_weight > 1.0 {
    1.0
  } else {
    semantic_weight
  }
  semantic_score * weight + lexical_score * (1.0 - weight)
}

///|
/// Rerank semantic results using document text overlap.
pub fn hybrid_rerank(
  query_text : String,
  semantic_results : Array[SearchResult],
  docs : Array[Document],
  semantic_weight : Double,
) -> Array[SearchResult] {
  let reranked = []
  for result in semantic_results {
    let lexical = match document_for_result(result, docs) {
      Some(doc) => lexical_overlap(query_text, doc)
      None => 0.0
    }
    reranked.push({
      id: result.id,
      score: hybrid_score(result.score, lexical, semantic_weight),
      metadata: result.metadata,
    })
  }
  sort_results(reranked, false)
  reranked
}

///|
/// Run a Flat semantic search followed by deterministic hybrid reranking.
pub fn hybrid_search(
  docs : Array[Document],
  query_vector : Array[Double],
  query_text : String,
  top_k : Int,
  semantic_weight : Double,
  filters : Array[(String, String)],
) -> Array[SearchResult] raise VectorError {
  let index = build_flat_index(docs)
  let candidate_k = if top_k < 1 { 1 } else { top_k * 3 }
  let semantic = index.search(query_vector, candidate_k, Cosine, filters)
  let reranked = hybrid_rerank(query_text, semantic, docs, semantic_weight)
  let limit = if top_k < reranked.length() { top_k } else { reranked.length() }
  let result = []
  for i = 0; i < limit; i = i + 1 {
    result.push(reranked[i])
  }
  result
}

///|
/// Build a prompt-ready RAG context in one call.
pub fn build_rag_context(
  docs : Array[Document],
  query_vector : Array[Double],
  query_text : String,
  top_k : Int,
  semantic_weight : Double,
  character_budget : Int,
) -> RetrievalContext raise VectorError {
  let results = hybrid_search(
    docs,
    query_vector,
    query_text,
    top_k,
    semantic_weight,
    [],
  )
  compose_retrieval_context(results, docs, character_budget, "\n\n")
}

///|
/// Return a compact source citation string for a context bundle.
pub fn context_citations(context : RetrievalContext) -> String {
  let citations = []
  for id in context.source_ids {
    citations.push("[" + id + "]")
  }
  citations.join(" ")
}

///|
fn metadata_first(metadata : Array[(String, String)], key : String) -> String {
  for pair in metadata {
    if pair.0 == key {
      return pair.1
    }
  }
  ""
}