import { tokenize } from "@/lib/text";
import type { Product } from "@/types/product";

/**
 * Keyword search: the lexical half of the hybrid ranker.
 *
 * BM25 over a single document per product, where the term frequency of each
 * field is scaled by that field's importance. A product matching in its name
 * therefore outranks an equally-matching product that only mentions the term
 * once in its description.
 */

export interface FieldSpec {
  name: string;
  weight: number;
  get: (product: Product) => string;
}

export const DEFAULT_FIELD_WEIGHTS: Record<string, number> = {
  name: 3,
  tags: 2.5,
  brand: 2,
  subcategory: 1.75,
  category: 1.5,
  description: 1,
};

export const DEFAULT_FIELDS: FieldSpec[] = [
  { name: "name", weight: DEFAULT_FIELD_WEIGHTS.name, get: (p) => p.name },
  { name: "tags", weight: DEFAULT_FIELD_WEIGHTS.tags, get: (p) => p.tags.join(" ") },
  { name: "brand", weight: DEFAULT_FIELD_WEIGHTS.brand, get: (p) => p.brand },
  {
    name: "subcategory",
    weight: DEFAULT_FIELD_WEIGHTS.subcategory,
    get: (p) => p.subcategory,
  },
  { name: "category", weight: DEFAULT_FIELD_WEIGHTS.category, get: (p) => p.category },
  {
    name: "description",
    weight: DEFAULT_FIELD_WEIGHTS.description,
    get: (p) => p.description,
  },
];

const K1 = 1.5;
const B = 0.75;

export interface KeywordMatch {
  productId: number;
  score: number;
  matchedTerms: string[];
}

interface IndexedProduct {
  product: Product;
  /** Weighted term frequency, already summed across fields. */
  frequencies: Map<string, number>;
  /** Weighted document length, used for BM25 length normalisation. */
  length: number;
}

function indexProduct(product: Product, fields: FieldSpec[]): IndexedProduct {
  const frequencies = new Map<string, number>();
  let length = 0;

  for (const field of fields) {
    for (const token of tokenize(field.get(product))) {
      const weighted = frequencies.get(token) ?? 0;
      const next = weighted + field.weight;
      frequencies.set(token, next);
      length += field.weight;
    }
  }

  return { product, frequencies, length };
}

/** Every stemmed token in the catalogue, used for query expansion. */
export function buildVocabulary(products: Product[], fields: FieldSpec[] = DEFAULT_FIELDS): Set<string> {
  const vocabulary = new Set<string>();
  for (const product of products) {
    for (const field of fields) {
      for (const token of tokenize(field.get(product))) vocabulary.add(token);
    }
  }
  return vocabulary;
}

/**
 * BM25 relevance for every product, normalised to 0..1 across the result set.
 * Products with no matching term are omitted.
 */
export function scoreKeyword(
  query: string,
  products: Product[],
  fields: FieldSpec[] = DEFAULT_FIELDS,
): KeywordMatch[] {
  const terms = [...new Set(tokenize(query))];
  if (terms.length === 0) return [];

  const indexed = products.map((product) => indexProduct(product, fields));
  const documentCount = indexed.length;
  if (documentCount === 0) return [];

  const documentFrequency = new Map<string, number>();
  for (const term of terms) {
    let frequency = 0;
    for (const entry of indexed) {
      if (entry.frequencies.has(term)) frequency += 1;
    }
    documentFrequency.set(term, frequency);
  }

  const averageLength = indexed.reduce((total, entry) => total + entry.length, 0) / documentCount;

  const matches: KeywordMatch[] = [];
  for (const entry of indexed) {
    let score = 0;
    const matchedTerms: string[] = [];

    for (const term of terms) {
      const termFrequency = entry.frequencies.get(term);
      if (termFrequency === undefined) continue;

      const frequency = documentFrequency.get(term) ?? 0;
      // BM25 idf, floored at a small positive value so that a term appearing in
      // every document still contributes rather than subtracting relevance.
      const idf = Math.max(0.05, Math.log(1 + (documentCount - frequency + 0.5) / (frequency + 0.5)));
      const denominator =
        termFrequency + K1 * (1 - B + (B * entry.length) / (averageLength || 1));
      score += idf * ((termFrequency * (K1 + 1)) / (denominator || 1));
      matchedTerms.push(term);
    }

    if (score > 0) {
      matches.push({ productId: entry.product.id, score, matchedTerms });
    }
  }

  const maxScore = matches.reduce((max, match) => Math.max(max, match.score), 0);
  if (maxScore <= 0) return matches;

  for (const match of matches) match.score = match.score / maxScore;
  return matches;
}