import type { MemoryEntry } from './entries.js';
import { scanProjectMemory } from './scan.js';
import { Bm25Ranker } from './bm25.js';
import { embed, isEmbeddingAvailable } from './embeddings.js';

const SEARCH_MODE = process.env['MEMORY_SEARCH'] ?? 'hybrid';

function extractKeywords(query: string): string[] {
  return query
    .toLowerCase()
    .replace(/[^a-z0-9\s]/g, '')
    .split(/\s+/)
    .filter((w) => w.length > 2);
}

export function rankByRelevance(query: string, entries: MemoryEntry[]): MemoryEntry[] {
  const keywords = extractKeywords(query);
  if (keywords.length === 0) return entries.slice(0, 5);

  const scored = entries.map((entry) => {
    let score = 0;
    const searchText = `${entry.title} ${entry.content} ${entry.tags.join(' ')}`.toLowerCase();
    for (const kw of keywords) {
      if (entry.title.toLowerCase().includes(kw)) score += 3;
      if (entry.tags.some((t) => t.toLowerCase().includes(kw))) score += 2;
      if (entry.content.toLowerCase().includes(kw)) score += 1;
    }
    return { entry, score };
  });

  return scored
    .filter((s) => s.score > 0)
    .sort((a, b) => b.score - a.score)
    .slice(0, 10)
    .map((s) => s.entry);
}

export async function rankByHybrid(
  query: string,
  entries: MemoryEntry[],
): Promise<MemoryEntry[]> {
  if (entries.length === 0) return [];
  const texts = entries.map((e) => `${e.title} ${e.content} ${e.tags.join(' ')}`);

  const bm25 = new Bm25Ranker();
  bm25.fit(texts);
  const bm25Scores = texts.map((_, i) => bm25.score(query, i));
  const maxBm25 = Math.max(...bm25Scores, 1);
  const normBm25 = bm25Scores.map((s) => s / maxBm25);

  let cosineScores: number[] = [];
  if (isEmbeddingAvailable()) {
    const vectors = await embed([query, ...texts]);
    if (vectors) {
      const queryVec = vectors[0]!;
      cosineScores = texts.map((_, i) => {
        const vec = vectors[i + 1];
        if (!vec) return 0;
        let dot = 0, nA = 0, nB = 0;
        for (let j = 0; j < queryVec.length; j++) {
          dot += queryVec[j]! * vec[j]!;
          nA += queryVec[j]! * queryVec[j]!;
          nB += vec[j]! * vec[j]!;
        }
        const denom = Math.sqrt(nA) * Math.sqrt(nB);
        return denom === 0 ? 0 : dot / denom;
      });
    }
  }

  const scored = entries.map((entry, i) => {
    const combined = (normBm25[i]! * 0.3) + ((cosineScores[i] ?? 0) * 0.7);
    return { entry, score: combined };
  });

  return scored
    .filter((s) => s.score >= 0.15)
    .sort((a, b) => b.score - a.score)
    .slice(0, 10)
    .map((s) => s.entry);
}

export async function loadMemoryForSession(
  projectRoot: string,
  _sessionId?: string,
  query?: string,
): Promise<string[]> {
  const entries = await scanProjectMemory(projectRoot);
  if (entries.length === 0) return [];

  const relevant = query
    ? SEARCH_MODE === 'keyword'
      ? rankByRelevance(query, entries)
      : await rankByHybrid(query, entries)
    : entries.slice(0, 5);

  return relevant.map((e) => `[${e.topic}] ${e.title}: ${e.content}`);
}

export { initEmbeddings } from './embeddings.js';