/**
 * Hybrid Search Implementation
 *
 * Combines vector (dense) search with BM25 (sparse) search for improved retrieval.
 * Supports multiple fusion methods: Reciprocal Rank Fusion (RRF) and Linear Combination.
 */
import { ProviderFactory } from "../../factories/providerFactory.js";
import { SpanSerializer, SpanType, SpanStatus, getMetricsAggregator, } from "../../observability/index.js";
import { logger } from "../../utils/logger.js";
import { rerank } from "../reranker/reranker.js";
/**
 * In-memory BM25 implementation for testing and development
 */
export class InMemoryBM25Index {
    documents = new Map();
    avgDocLength = 0;
    k1 = 1.5; // BM25 parameter
    b = 0.75; // BM25 parameter
    async search(query, topK = 10) {
        const queryTokens = this.tokenize(query);
        if (queryTokens.length === 0 || this.documents.size === 0) {
            return [];
        }
        // Calculate IDF for each query term
        const idfValues = new Map();
        for (const token of queryTokens) {
            const docCount = this.countDocumentsWithTerm(token);
            const idf = Math.log((this.documents.size - docCount + 0.5) / (docCount + 0.5) + 1);
            idfValues.set(token, idf);
        }
        // Calculate BM25 score for each document
        const scores = [];
        for (const [id, doc] of this.documents) {
            let score = 0;
            const docLength = doc.tokens.length;
            for (const token of queryTokens) {
                const tf = this.countTermFrequency(doc.tokens, token);
                const idf = idfValues.get(token) || 0;
                // BM25 scoring formula
                const numerator = tf * (this.k1 + 1);
                const denominator = tf +
                    this.k1 * (1 - this.b + this.b * (docLength / this.avgDocLength));
                score += idf * (numerator / denominator);
            }
            if (score > 0) {
                scores.push({
                    id,
                    score,
                    text: doc.text,
                    metadata: doc.metadata,
                });
            }
        }
        // Sort by score descending
        scores.sort((a, b) => b.score - a.score);
        return scores.slice(0, topK);
    }
    async addDocuments(documents) {
        for (const doc of documents) {
            const tokens = this.tokenize(doc.text);
            this.documents.set(doc.id, {
                text: doc.text,
                tokens,
                metadata: doc.metadata || {},
            });
        }
        // Recalculate average document length
        let totalLength = 0;
        for (const doc of this.documents.values()) {
            totalLength += doc.tokens.length;
        }
        this.avgDocLength =
            this.documents.size > 0 ? totalLength / this.documents.size : 0;
    }
    tokenize(text) {
        return text
            .toLowerCase()
            .replace(/[^\w\s]/g, " ")
            .split(/\s+/)
            .filter((t) => t.length > 0);
    }
    countTermFrequency(tokens, term) {
        return tokens.filter((t) => t === term).length;
    }
    countDocumentsWithTerm(term) {
        let count = 0;
        for (const doc of this.documents.values()) {
            if (doc.tokens.includes(term)) {
                count++;
            }
        }
        return count;
    }
}
/**
 * Reciprocal Rank Fusion
 * Combines rankings from multiple retrieval methods
 *
 * @param rankings - Array of ranking lists, each with id and rank
 * @param k - RRF constant (default: 60)
 * @returns Map of document IDs to fused scores
 */
export function reciprocalRankFusion(rankings, k = 60) {
    const scores = new Map();
    for (const ranking of rankings) {
        for (const { id, rank } of ranking) {
            const currentScore = scores.get(id) || 0;
            scores.set(id, currentScore + 1 / (k + rank));
        }
    }
    return scores;
}
/**
 * Linear Combination of normalized scores
 *
 * @param vectorScores - Vector search scores
 * @param bm25Scores - BM25 search scores
 * @param alpha - Weight for vector scores (0-1), bm25 gets 1-alpha
 * @returns Map of document IDs to combined scores
 */
export function linearCombination(vectorScores, bm25Scores, alpha = 0.5) {
    const combined = new Map();
    // Get all document IDs
    const allIds = new Set([...vectorScores.keys(), ...bm25Scores.keys()]);
    // Normalize scores
    const normalizedVector = normalizeScores(vectorScores);
    const normalizedBM25 = normalizeScores(bm25Scores);
    for (const id of allIds) {
        const vectorScore = normalizedVector.get(id) || 0;
        const bm25Score = normalizedBM25.get(id) || 0;
        combined.set(id, alpha * vectorScore + (1 - alpha) * bm25Score);
    }
    return combined;
}
/**
 * Normalize scores to 0-1 range
 */
function normalizeScores(scores) {
    const values = Array.from(scores.values());
    if (values.length === 0) {
        return new Map();
    }
    const min = Math.min(...values);
    const max = Math.max(...values);
    const range = max - min || 1;
    const normalized = new Map();
    for (const [id, score] of scores) {
        normalized.set(id, (score - min) / range);
    }
    return normalized;
}
/**
 * Create a hybrid search function
 *
 * @param options - Search options
 * @returns Hybrid search function
 */
export function createHybridSearch(options) {
    const { vectorStore, bm25Index, indexName, embeddingModel, defaultConfig = {}, } = options;
    /**
     * Execute hybrid search combining vector and BM25 retrieval
     *
     * @param query - Search query
     * @param config - Search configuration
     * @returns Hybrid search results
     */
    return async function hybridSearch(query, config) {
        const startTime = Date.now();
        const { vectorWeight = defaultConfig.vectorWeight ?? 0.5, bm25Weight = defaultConfig.bm25Weight ?? 0.5, fusionMethod = defaultConfig.fusionMethod ?? "rrf", rrfK = defaultConfig.rrfK ?? 60, topK = defaultConfig.topK ?? 10, enableReranking = defaultConfig.enableReranking ?? false, reranker: rerankerConfig = defaultConfig.reranker, } = config || {};
        const span = SpanSerializer.createSpan(SpanType.RAG, "rag.search", {
            "rag.operation": "search",
            "rag.topK": topK,
            "rag.fusionMethod": fusionMethod,
            "rag.query": query.slice(0, 200),
        });
        const spanStartTime = Date.now();
        try {
            // Generate query embedding
            const embeddingProvider = await ProviderFactory.createProvider(embeddingModel?.provider, embeddingModel?.modelName);
            if (typeof embeddingProvider
                .embed !== "function") {
                throw new Error(`Embedding provider does not support the embed() method. ` +
                    `Please use a provider that supports embeddings (e.g., OpenAI text-embedding-3-small, Vertex text-embedding-004).`);
            }
            const queryEmbedding = await embeddingProvider.embed(query);
            // Parallel retrieval
            const [vectorResults, bm25Results] = await Promise.all([
                vectorStore.query({
                    indexName,
                    queryVector: queryEmbedding,
                    topK: topK * 2, // Get more for fusion
                }),
                bm25Index.search(query, topK * 2),
            ]);
            // Fuse results
            let fusedResults;
            if (fusionMethod === "rrf") {
                // Reciprocal Rank Fusion
                const vectorRanking = vectorResults.map((r, i) => ({
                    id: r.id,
                    rank: i + 1,
                }));
                const bm25Ranking = bm25Results.map((r, i) => ({
                    id: r.id,
                    rank: i + 1,
                }));
                const rrfScores = reciprocalRankFusion([vectorRanking, bm25Ranking], rrfK);
                // Combine with original data
                const resultMap = new Map();
                for (const r of vectorResults) {
                    resultMap.set(r.id, { text: r.text || "", metadata: r.metadata });
                }
                for (const r of bm25Results) {
                    if (!resultMap.has(r.id)) {
                        resultMap.set(r.id, { text: r.text, metadata: r.metadata });
                    }
                }
                fusedResults = Array.from(rrfScores.entries())
                    .sort((a, b) => b[1] - a[1])
                    .slice(0, topK)
                    .map(([id, score]) => ({
                    id,
                    score,
                    text: resultMap.get(id)?.text || "",
                    metadata: resultMap.get(id)?.metadata,
                    scores: {
                        combined: score,
                    },
                }));
            }
            else {
                // Linear combination
                const vectorScoreMap = new Map(vectorResults.map((r) => [r.id, r.score || 0]));
                const bm25ScoreMap = new Map(bm25Results.map((r) => [r.id, r.score]));
                // Adjust weights based on config
                const totalWeight = vectorWeight + bm25Weight;
                const normalizedVectorWeight = vectorWeight / totalWeight;
                const combinedScores = linearCombination(vectorScoreMap, bm25ScoreMap, normalizedVectorWeight);
                // Combine with original data
                const resultMap = new Map();
                for (const r of vectorResults) {
                    resultMap.set(r.id, {
                        text: r.text || "",
                        metadata: r.metadata,
                        vectorScore: r.score,
                    });
                }
                for (const r of bm25Results) {
                    const existing = resultMap.get(r.id);
                    if (existing) {
                        existing.bm25Score = r.score;
                    }
                    else {
                        resultMap.set(r.id, {
                            text: r.text,
                            metadata: r.metadata,
                            bm25Score: r.score,
                        });
                    }
                }
                fusedResults = Array.from(combinedScores.entries())
                    .sort((a, b) => b[1] - a[1])
                    .slice(0, topK)
                    .map(([id, score]) => {
                    const data = resultMap.get(id);
                    return {
                        id,
                        score,
                        text: data?.text || "",
                        metadata: data?.metadata,
                        scores: {
                            vector: data?.vectorScore,
                            bm25: data?.bm25Score,
                            combined: score,
                        },
                    };
                });
            }
            // Apply reranking if configured
            if (enableReranking && rerankerConfig && fusedResults.length > 0) {
                const rerankerModel = await ProviderFactory.createProvider(rerankerConfig.model.provider, rerankerConfig.model.modelName);
                const rerankedResults = await rerank(fusedResults.map((r) => ({
                    id: r.id,
                    text: r.text,
                    score: r.score,
                    metadata: r.metadata,
                })), query, rerankerModel, {
                    weights: rerankerConfig.weights,
                    topK: rerankerConfig.topK || topK,
                });
                fusedResults = rerankedResults.map((r) => ({
                    id: r.result.id,
                    score: r.score,
                    text: r.result.text || "",
                    metadata: r.result.metadata,
                    scores: {
                        ...(fusedResults.find((f) => f.id === r.result.id)?.scores || {}),
                        reranked: r.score,
                    },
                }));
            }
            const queryTime = Date.now() - startTime;
            logger.info("[HybridSearch] Search completed", {
                query: query.slice(0, 50),
                vectorResults: vectorResults.length,
                bm25Results: bm25Results.length,
                fusedResults: fusedResults.length,
                fusionMethod,
                queryTime,
            });
            span.durationMs = Date.now() - spanStartTime;
            const endedSpan = SpanSerializer.endSpan(span, SpanStatus.OK);
            endedSpan.attributes = {
                ...endedSpan.attributes,
                "rag.results_count": fusedResults.length,
                "rag.vector_results": vectorResults.length,
                "rag.bm25_results": bm25Results.length,
            };
            getMetricsAggregator().recordSpan(endedSpan);
            return fusedResults;
        }
        catch (error) {
            span.durationMs = Date.now() - spanStartTime;
            const endedSpan = SpanSerializer.endSpan(span, SpanStatus.ERROR);
            endedSpan.statusMessage =
                error instanceof Error ? error.message : String(error);
            getMetricsAggregator().recordSpan(endedSpan);
            logger.error("[HybridSearch] Search failed", {
                query: query.slice(0, 50),
                error: error instanceof Error ? error.message : String(error),
            });
            throw error;
        }
    };
}
