import type { Card } from "@/lib/types";

export type PromptModel = { version: 1; dimensions: 50; words: Record<string, number[]> };
export type PromptMatch = { card: Card; score: number; keywordScore: number; semanticScore: number; method: "keyword" | "semantic" };
// Frozen before evaluation. Similarity is a ranking signal, never a probability.
export const SEMANTIC_FLOOR = 0.55;
const stopWords = new Set("a an and are as at be been but by can could for from had has have how i if in into is it its like me my of on or our so some than that the their them then there these they this to us was we were what when where which who with would you your play song track music".split(" "));

export function promptTokens(text: string) {
  return (text.toLowerCase().match(/[a-z]+/g) || []).filter((word) => !stopWords.has(word));
}

export function validatePromptModel(value: unknown): PromptModel {
  const model = value as PromptModel | null;
  if (!model || model.version !== 1 || model.dimensions !== 50 || !model.words || typeof model.words !== "object") throw new Error("Invalid prompt model.");
  const entries = Object.entries(model.words);
  if (!entries.length || entries.length > 20000 || entries.some(([word, vector]) =>
    !/^[a-z]+$/.test(word) || !Array.isArray(vector) || vector.length !== 50 || vector.some((x) => !Number.isInteger(x) || Math.abs(x) > 127))) {
    throw new Error("Invalid prompt vectors.");
  }
  return model;
}

function vectorFor(tokens: string[], model: PromptModel, frequency: Map<string, number>) {
  const vector = Array<number>(50).fill(0);
  for (const word of tokens) {
    if (!Object.hasOwn(model.words, word)) continue;
    const source = model.words[word];
    const weight = 1 / Math.sqrt(1 + (frequency.get(word) || 0));
    for (let i = 0; i < 50; i++) vector[i] += source[i] * weight;
  }
  const norm = Math.hypot(...vector);
  return norm ? vector.map((x) => x / norm) : vector;
}

export function searchPrompts(query: string, cards: Card[], options: {
  model?: PromptModel; category?: string; excludedIds?: string[]; limit?: number;
} = {}): PromptMatch[] {
  if (!query.trim() || query.length > 240 || !cards.length) return [];
  const terms = [...new Set(promptTokens(query))];
  if (!terms.length) return [];
  const documents = cards.map((card) => promptTokens(`${card.category} ${card.prompt}`));
  const frequency = new Map<string, number>();
  for (const doc of documents) for (const token of new Set(doc)) frequency.set(token, (frequency.get(token) || 0) + 1);
  const avgLength = documents.reduce((sum, doc) => sum + doc.length, 0) / cards.length || 1;
  const queryVector = options.model ? vectorFor(terms, options.model, frequency) : null;
  const excluded = new Set(options.excludedIds);
  return cards.flatMap((card, index): PromptMatch[] => {
    if (excluded.has(card.id) || (options.category && card.category !== options.category)) return [];
    const doc = documents[index];
    const keywordScore = terms.reduce((sum, term) => {
      const tf = doc.filter((word) => word === term).length;
      const df = frequency.get(term) || 0;
      const idf = Math.log(1 + (cards.length - df + 0.5) / (df + 0.5));
      return sum + idf * (tf * 2.2) / (tf + 1.2 * (0.25 + 0.75 * doc.length / avgLength));
    }, 0);
    const semanticScore = options.model && queryVector
      ? vectorFor(doc, options.model, frequency).reduce((sum, x, i) => sum + x * queryVector[i], 0) : 0;
    if (keywordScore <= 0 && (!options.model || semanticScore < SEMANTIC_FLOOR)) return [];
    const score = options.model ? 0.65 * semanticScore + 0.35 * keywordScore / (keywordScore + 3) : keywordScore;
    return [{ card, score, keywordScore, semanticScore, method: options.model ? "semantic" : "keyword" }];
  }).sort((a, b) => b.score - a.score || a.card.id.localeCompare(b.card.id)).slice(0, options.limit ?? 5);
}