/** Field-aware BM25 retrieval over immutable tool documents. @module dsh-tool-retriever/retrieval/bm25f */

import type { ResolvedConfig } from '../config.ts'
import { normalizeIdentifier, tokenize } from './tokenizer.ts'
import type { ToolDocument, ToolDocumentField, ToolSearchHit } from './types.ts'

const FIELDS: readonly ToolDocumentField[] = ['name', 'description', 'parameterName', 'parameterDescription']

/**
 * Rank one live catalog with strict canonical-name matching followed by BM25F.
 * @param query - Raw model-supplied query.
 * @param documents - Deterministic projections of currently visible tools.
 * @param config - Resolved scoring and result-count configuration.
 * @returns Stable candidates ordered by exact match, BM25F score, then name.
 */
export function retrieveTools(query: string, documents: readonly ToolDocument[], config: ResolvedConfig): ToolSearchHit[] {
  const queryTokens = [...new Set(tokenize(query))]
  const normalizedQuery = normalizeIdentifier(query)
  const averageLengths = averageFieldLengths(documents)
  const documentFrequency = termDocumentFrequency(documents, queryTokens)
  const hits = documents.flatMap((document): ToolSearchHit[] => {
    const exact = normalizedQuery !== '' && normalizedQuery === document.normalizedName
    const matchedFields = FIELDS.filter(field => intersects(document.fields[field], queryTokens))
    const bm25 = scoreDocument(document, queryTokens, documentFrequency, averageLengths, documents.length, config)
    if (!exact && bm25 <= 0) return []
    return [{
      name: document.name,
      description: document.description,
      rank: 0,
      score: bm25,
      exact,
      bm25,
      matchedFields,
    }]
  })
  hits.sort((left, right) => Number(right.exact) - Number(left.exact)
    || right.bm25 - left.bm25
    || left.name.localeCompare(right.name))
  return hits.slice(0, config.topK).map((hit, index) => ({ ...hit, rank: index + 1 }))
}

function averageFieldLengths(documents: readonly ToolDocument[]): Record<ToolDocumentField, number> {
  const count = Math.max(documents.length, 1)
  return Object.fromEntries(FIELDS.map(field => [
    field,
    documents.reduce((sum, document) => sum + document.fields[field].length, 0) / count,
  ])) as Record<ToolDocumentField, number>
}

function termDocumentFrequency(documents: readonly ToolDocument[], terms: readonly string[]): Map<string, number> {
  const frequencies = new Map<string, number>()
  for (const term of terms) {
    let count = 0
    for (const document of documents) {
      if (FIELDS.some(field => document.fields[field].includes(term))) count += 1
    }
    frequencies.set(term, count)
  }
  return frequencies
}

function scoreDocument(
  document: ToolDocument,
  terms: readonly string[],
  documentFrequency: ReadonlyMap<string, number>,
  averageLengths: Readonly<Record<ToolDocumentField, number>>,
  documentCount: number,
  config: ResolvedConfig,
): number {
  let score = 0
  for (const term of terms) {
    let weightedFrequency = 0
    for (const field of FIELDS) {
      const values = document.fields[field]
      const frequency = values.reduce((count, token) => count + Number(token === term), 0)
      if (frequency === 0) continue
      const averageLength = averageLengths[field]
      const lengthRatio = averageLength === 0 ? 1 : values.length / averageLength
      const normalization = 1 - config.bm25B + config.bm25B * lengthRatio
      weightedFrequency += config.fieldWeights[field] * frequency / normalization
    }
    if (weightedFrequency === 0) continue
    const frequency = documentFrequency.get(term) ?? 0
    const inverseDocumentFrequency = Math.log(1 + (documentCount - frequency + 0.5) / (frequency + 0.5))
    score += inverseDocumentFrequency
      * weightedFrequency * (config.bm25K1 + 1)
      / (weightedFrequency + config.bm25K1)
  }
  return score
}

function intersects(tokens: readonly string[], query: readonly string[]): boolean {
  return query.some(term => tokens.includes(term))
}