import { tokenize } from "./text";
import type { PortfolioDocument, SearchHit, SearchOptions } from "./types";

interface Field {
  weight: number;
  get: (d: PortfolioDocument) => string;
}

const FIELDS: Field[] = [
  { weight: 3.2, get: (d) => d.title },
  { weight: 2.2, get: (d) => (d.tags || []).join(" ") },
  { weight: 1.4, get: (d) => d.summary },
  { weight: 1.0, get: (d) => (d.facts || []).join(" ") },
  { weight: 0.6, get: (d) => d.body || "" },
  { weight: 0.8, get: (d) => d.category },
];

export interface SearchIndex {
  docs: PortfolioDocument[];
  /** term → docIndex → weighted term frequency */
  postings: Map<string, Map<number, number>>;
  lengths: number[];
  avgLength: number;
  /** prefix lookup for instant (as-you-type) search */
  terms: string[];
}

/** Build an in-memory BM25F index. ~1ms for a few hundred documents. */
export function createIndex(docs: PortfolioDocument[]): SearchIndex {
  const postings = new Map<string, Map<number, number>>();
  const lengths: number[] = [];
  docs.forEach((doc, i) => {
    let len = 0;
    for (const f of FIELDS) {
      const toks = tokenize(f.get(doc));
      len += toks.length * f.weight;
      for (const t of toks) {
        let m = postings.get(t);
        if (!m) postings.set(t, (m = new Map()));
        m.set(i, (m.get(i) || 0) + f.weight);
      }
    }
    lengths.push(len || 1);
  });
  const avgLength = lengths.reduce((a, b) => a + b, 0) / Math.max(1, lengths.length);
  return { docs, postings, lengths, avgLength, terms: [...postings.keys()].sort() };
}

function expand(index: SearchIndex, term: string, isLast: boolean): { term: string; factor: number }[] {
  const out = [{ term, factor: 1 }];
  if (isLast && term.length >= 2) {
    // prefix expansion for the token being typed
    for (const t of index.terms) {
      if (t !== term && t.startsWith(term)) out.push({ term: t, factor: 0.7 });
      if (out.length > 12) break;
    }
  }
  return out;
}

export function search(index: SearchIndex, query: string, opts: SearchOptions = {}): SearchHit[] {
  const k1 = 1.2;
  const b = 0.72;
  const N = index.docs.length;
  const qTerms = [...new Set(tokenize(query))];
  if (!qTerms.length) return [];
  const scores = new Map<number, { score: number; matched: Set<string> }>();
  qTerms.forEach((qt, qi) => {
    for (const { term, factor } of expand(index, qt, qi === qTerms.length - 1)) {
      const post = index.postings.get(term);
      if (!post) continue;
      const idf = Math.log(1 + (N - post.size + 0.5) / (post.size + 0.5));
      for (const [di, tf] of post) {
        const norm = tf + k1 * (1 - b + (b * index.lengths[di]) / index.avgLength);
        const s = idf * ((tf * (k1 + 1)) / norm) * factor;
        const cur = scores.get(di) || { score: 0, matched: new Set<string>() };
        cur.score += s;
        cur.matched.add(qt);
        scores.set(di, cur);
      }
    }
  });
  const hits: SearchHit[] = [];
  for (const [di, { score, matched }] of scores) {
    const doc = index.docs[di];
    if (opts.categories && !opts.categories.includes(doc.category)) continue;
    // coverage bonus: prefer docs matching more distinct query terms
    const coverage = matched.size / qTerms.length;
    const s = score * (0.55 + 0.45 * coverage) * (doc.weight ?? 1) * (opts.boost ? opts.boost(doc) : 1);
    hits.push({ doc, score: s, matched: [...matched] });
  }
  hits.sort((a, b2) => b2.score - a.score);
  return hits.slice(0, opts.limit ?? 10);
}