/**
 * RAG Retrieval Module
 * Generated by vai v{{vaiVersion}} on {{generatedAt}}
 * 
 * Model: {{model}}
 * Database: {{db}}.{{collection}}
 * Reranking: {{#if rerank}}enabled ({{rerankModel}}){{else}}disabled{{/if}}
 */

const { embed{{#if rerank}}, rerank{{/if}} } = require('./client');
const { vectorSearch } = require('./connection');

/**
 * Retrieve relevant documents for a query.
 * 
 * Pipeline:
 * 1. Embed the query using Voyage AI
 * 2. Vector search in MongoDB Atlas
{{#if rerank}}
 * 3. Rerank results using Voyage AI
{{/if}}
 * 
 * @param {string} query - User's query
 * @param {object} options - Retrieval options
 * @param {number} options.limit - Final number of results (default: 5)
 * @param {number} options.candidates - Initial candidates for reranking (default: 20)
 * @param {object} options.filter - MongoDB pre-filter
 * @returns {Promise<Array<{text: string, score: number, metadata: object}>>}
 */
async function retrieve(query, options = {}) {
  const limit = options.limit || 5;
{{#if rerank}}
  const candidates = options.candidates || 20;
{{/if}}

  // Step 1: Embed the query
  const { embeddings } = await embed(query, { inputType: 'query' });
  const queryEmbedding = embeddings[0];

  // Step 2: Vector search
{{#if rerank}}
  const searchResults = await vectorSearch(queryEmbedding, {
    limit: candidates,
    filter: options.filter,
  });
{{else}}
  const searchResults = await vectorSearch(queryEmbedding, {
    limit,
    filter: options.filter,
  });
{{/if}}

  if (searchResults.length === 0) {
    return [];
  }

{{#if rerank}}
  // Step 3: Rerank for better relevance
  const documents = searchResults.map(r => r.document.text);
  const { results: reranked } = await rerank(query, documents, { topK: limit });

  // Map reranked results back to full documents
  return reranked.map(r => ({
    text: searchResults[r.index].document.text,
    score: r.relevanceScore,
    metadata: searchResults[r.index].document.metadata,
  }));
{{else}}
  return searchResults.map(r => ({
    text: r.document.text,
    score: r.score,
    metadata: r.document.metadata,
  }));
{{/if}}
}

/**
 * Format retrieved documents as context for an LLM prompt.
 * @param {Array<{text: string, score: number}>} results - Retrieved documents
 * @param {object} options - Formatting options
 * @param {boolean} options.includeScores - Include relevance scores (default: false)
 * @param {string} options.separator - Separator between documents (default: '\n\n---\n\n')
 * @returns {string} Formatted context string
 */
function formatContext(results, options = {}) {
  const separator = options.separator || '\n\n---\n\n';
  
  return results.map((r, i) => {
    let text = r.text;
    if (options.includeScores) {
      text = `[Score: ${r.score.toFixed(3)}]\n${text}`;
    }
    return text;
  }).join(separator);
}

module.exports = {
  retrieve,
  formatContext,
};