import type { ModelCapabilities, PolyModel, ProviderId } from "@/lib/ai/core/types";
import { blendedPrice, buildBadgeContext, isCodingModel, isFastModel, isFrontierModel, isNewModel, isOpenWeightsModel, releaseDateMs, firstSeenMs, sortWeightOf, type BadgeContext } from "./badges";

/**
 * Model search with intent parsing — pure, no React.
 *
 *   "cheap vision model"      → capability vision, sort cheapest
 *   "1M context"              → minContext 1 000 000
 *   "under $1/M"              → maxInputPrice 1
 *   "fastest gemini"          → brand gemini, sort fastest
 *   "json schema" / "200k" / "open source" / "reasoning" / "OpenAI"
 *
 * Free-text tokens must all match (name, id, family, provider, user label); intents filter or boost.
 */
export type CapabilityKey = keyof ModelCapabilities;
export type SearchSort = "cheapest" | "fastest" | "context" | "newest" | "quality";

export interface SearchIntent {
  /** Remaining free-text tokens (already lower-cased). */
  text: string[];
  brands: Brand[];
  capabilities: CapabilityKey[];
  openWeights: boolean;
  coding: boolean;
  newest: boolean;
  quality: boolean;
  minContext?: number;
  maxInputPrice?: number;
  maxOutputPrice?: number;
  sort?: SearchSort;
  /** Human-readable chips describing what was understood. */
  chips: string[];
}

export interface Brand {
  key: string;
  label: string;
  provider?: ProviderId;
  match: RegExp;
}

const BRANDS: (Brand & { words: RegExp })[] = [
  { key: "openai", label: "OpenAI", provider: "openai", words: /^(openai|gpt|chatgpt|o[1-4])$/, match: /\bgpt|openai|\bo[1-4]\b/i },
  { key: "anthropic", label: "Anthropic", provider: "anthropic", words: /^(anthropic|claude)$/, match: /claude|anthropic/i },
  { key: "gemini", label: "Gemini", provider: "gemini", words: /^(gemini|google)$/, match: /gemini|google/i },
  { key: "xai", label: "xAI", provider: "xai", words: /^(xai|grok)$/, match: /grok|x-ai|\bxai\b/i },
  { key: "mistral", label: "Mistral", provider: "mistral", words: /^(mistral|mistralai)$/, match: /mistral|codestral|magistral|ministral|devstral|mixtral/i },
  { key: "deepseek", label: "DeepSeek", provider: "deepseek", words: /^deepseek$/, match: /deepseek/i },
  { key: "kimi", label: "Kimi", provider: "kimi", words: /^(kimi|moonshot|moonshotai)$/, match: /kimi|moonshot/i },
  { key: "openrouter", label: "OpenRouter", provider: "openrouter", words: /^openrouter$/, match: /^$/ },
  { key: "cerebras", label: "Cerebras", provider: "cerebras", words: /^cerebras$/, match: /^$/ },
  { key: "llama", label: "Llama", words: /^(llama|meta)$/, match: /llama|meta-/i },
  { key: "qwen", label: "Qwen", words: /^qwen$/, match: /qwen/i },
];

const CAP_WORDS: { re: RegExp; cap: CapabilityKey; chip: string }[] = [
  { re: /^(vision|image|images|multimodal|photo|photos|ocr|see|picture|pictures)$/, cap: "vision", chip: "Vision" },
  { re: /^(reasoning|reason|reasoner|thinking|think|thinks|cot|deliberate)$/, cap: "reasoning", chip: "Reasoning" },
  { re: /^(tool|tools|function|functions|agentic|agent|agents)$/, cap: "tools", chip: "Tools" },
  { re: /^(web|search|browse|browsing|internet|grounding|grounded)$/, cap: "webSearch", chip: "Web search" },
  { re: /^(json|schema|structured)$/, cap: "structuredOutput", chip: "JSON schema" },
  { re: /^(pdf|pdfs|file|files|document|documents)$/, cap: "files", chip: "PDF / files" },
  { re: /^(audio|voice|speech)$/, cap: "audioInput", chip: "Audio input" },
  { re: /^(video)$/, cap: "video", chip: "Video" },
];

const PHRASES: [RegExp, string][] = [
  [/open[\s-]?(source|weights?)/g, " openweights "],
  [/\boss\b/g, " openweights "],
  [/json[\s-]?schema/g, " json "],
  [/structured[\s-]?output(s)?/g, " json "],
  [/web[\s-]?search/g, " web "],
  [/(tool|function)[\s-]?calling/g, " tools "],
  [/(long|large|big|huge)[\s-]?(context|ctx|window)/g, " longcontext "],
  [/context[\s-]?window/g, " context "],
  [/low[\s-]?cost/g, " cheap "],
  [/low[\s-]?latency/g, " fast "],
  [/per[\s-]?million/g, "/m"],
  [/\bmtok\b/g, "/m"],
];

const CHEAP_RE = /^(cheap|cheapest|cheaper|budget|inexpensive|affordable|economical|free)$/;
const FAST_WORD_RE = /^(fast|fastest|faster|quick|quickest|speedy|speed|snappy|instant|latency)$/;
const OPEN_RE = /^openweights$/;
const CODING_WORD_RE = /^(code|coding|coder|programming|program|developer|dev|debug|refactor)$/;
const NEW_RE = /^(new|newest|latest|recent|recently|released|fresh)$/;
const QUALITY_RE = /^(best|smart|smartest|strongest|frontier|quality|powerful|top|flagship|premium)$/;
const LONG_RE = /^(longcontext|long)$/;
const STOP_RE = /^(a|an|the|model|models|llm|llms|with|for|and|or|that|which|to|in|of|is|are|can|good|great|at|my|me|i|want|need|please|show|find)$/;

function ctxTokens(n: number, unit: string): number {
  return unit === "m" ? Math.round(n * 1_000_000) : Math.round(n * 1_000);
}

function fmtCtx(n: number): string {
  return n >= 1_000_000 ? `${(n / 1_000_000).toFixed(n % 1_000_000 ? 1 : 0)}M` : `${Math.round(n / 1000)}K`;
}

export function parseSearchQuery(raw: string): SearchIntent {
  let q = ` ${raw.toLowerCase().trim()} `;
  for (const [re, rep] of PHRASES) q = q.replace(re, rep);
  const intent: SearchIntent = { text: [], brands: [], capabilities: [], openWeights: false, coding: false, newest: false, quality: false, chips: [] };

  // Price constraints: "under $1/m", "< $2", "$0.5 per million", "2$/m". Output when the query says so.
  const wantsOutput = /\b(output|out|completion)\b/.test(q);
  const priceRe = /(?:(under|below|<|less than|max|up to|at most|cheaper than)\s*)?\$\s?(\d+(?:\.\d+)?)(?:\s*\/\s*1?m)?|(?:(under|below|<|less than|max|up to|at most|cheaper than)\s*)?(\d+(?:\.\d+)?)\s?\$(?:\s*\/\s*1?m)?|(\d+(?:\.\d+)?)\s*\/\s*1?m\b/g;
  q = q.replace(priceRe, (_m, _p1, a, _p3, b, c) => {
    const v = Number(a ?? b ?? c);
    if (Number.isFinite(v)) {
      if (wantsOutput) intent.maxOutputPrice = v;
      else intent.maxInputPrice = v;
    }
    return " ";
  });
  q = q.replace(/\b(output|out|completion|input|in)\b/g, " ");

  // Context: "1m context", "200k", "128k ctx", "1000000 tokens"
  const ctxRe = /(\d+(?:\.\d+)?)\s*([mk])\b(?:\s*(?:context|ctx|tokens?|window))?/g;
  q = q.replace(ctxRe, (_m, n, unit) => {
    const v = ctxTokens(Number(n), unit);
    intent.minContext = Math.max(intent.minContext ?? 0, v);
    return " ";
  });
  q = q.replace(/(\d{5,8})\s*(?:context|ctx|tokens?|window)/g, (_m, n) => {
    intent.minContext = Math.max(intent.minContext ?? 0, Number(n));
    return " ";
  });
  q = q.replace(/\b(context|ctx|tokens?|window)\b/g, " ");

  for (const tok of q.split(/\s+/).filter(Boolean)) {
    const brand = BRANDS.find((b) => b.words.test(tok));
    if (brand) {
      if (!intent.brands.some((b) => b.key === brand.key)) intent.brands.push(brand);
      continue;
    }
    const cap = CAP_WORDS.find((c) => c.re.test(tok));
    if (cap) {
      if (!intent.capabilities.includes(cap.cap)) intent.capabilities.push(cap.cap);
      continue;
    }
    if (CHEAP_RE.test(tok)) {
      intent.sort = intent.sort ?? "cheapest";
      continue;
    }
    if (FAST_WORD_RE.test(tok)) {
      intent.sort = intent.sort ?? "fastest";
      continue;
    }
    if (OPEN_RE.test(tok)) {
      intent.openWeights = true;
      continue;
    }
    if (CODING_WORD_RE.test(tok)) {
      intent.coding = true;
      continue;
    }
    if (NEW_RE.test(tok)) {
      intent.newest = true;
      intent.sort = intent.sort ?? "newest";
      continue;
    }
    if (QUALITY_RE.test(tok)) {
      intent.quality = true;
      intent.sort = intent.sort ?? "quality";
      continue;
    }
    if (LONG_RE.test(tok)) {
      intent.minContext = Math.max(intent.minContext ?? 0, 400_000);
      intent.sort = intent.sort ?? "context";
      continue;
    }
    if (STOP_RE.test(tok)) continue;
    intent.text.push(tok);
  }

  for (const b of intent.brands) intent.chips.push(b.label);
  for (const c of intent.capabilities) intent.chips.push(CAP_WORDS.find((w) => w.cap === c)?.chip ?? c);
  if (intent.minContext) intent.chips.push(`≥ ${fmtCtx(intent.minContext)} context`);
  if (intent.maxInputPrice !== undefined) intent.chips.push(`≤ $${intent.maxInputPrice}/M input`);
  if (intent.maxOutputPrice !== undefined) intent.chips.push(`≤ $${intent.maxOutputPrice}/M output`);
  if (intent.openWeights) intent.chips.push("Open weights");
  if (intent.coding) intent.chips.push("Coding");
  if (intent.newest) intent.chips.push("New");
  if (intent.sort === "cheapest") intent.chips.push("Cheapest first");
  if (intent.sort === "fastest") intent.chips.push("Fastest first");
  if (intent.sort === "quality") intent.chips.push("Strongest first");
  return intent;
}

export function isEmptyIntent(i: SearchIntent): boolean {
  return !i.text.length && !i.brands.length && !i.capabilities.length && !i.openWeights && !i.coding && !i.newest && !i.quality && i.minContext === undefined && i.maxInputPrice === undefined && i.maxOutputPrice === undefined && !i.sort;
}

export interface SearchContext {
  ctx?: BadgeContext;
  favorites?: Set<string>;
  recents?: Set<string> | string[];
  labels?: Record<string, string>;
  providerNames?: Partial<Record<ProviderId, string>>;
  now?: number;
}

export interface SearchResult {
  model: PolyModel;
  score: number;
  /** Why it matched (for the UI). */
  matches: string[];
}

/** Subsequence match ("sonet" ⊂ "sonnet") for light typo tolerance. */
function fuzzy(needle: string, hay: string): boolean {
  let i = 0;
  for (const ch of hay) {
    if (ch === needle[i]) i++;
    if (i === needle.length) return true;
  }
  return i === needle.length;
}

function textScore(m: PolyModel, tokens: string[], whole: string, label: string | undefined, providerName: string): number {
  if (!tokens.length) return 0;
  const name = m.displayName.toLowerCase();
  const id = m.id.toLowerCase();
  const family = (m.family ?? "").toLowerCase();
  const vendor = String(m.metadata?.vendor ?? "").toLowerCase();
  const lbl = (label ?? "").toLowerCase();
  const prov = providerName.toLowerCase();
  let score = 0;
  if (id === whole || name === whole) score += 100;
  else if (name.startsWith(whole) || id.startsWith(whole)) score += 60;
  const nameWords = name.split(/[\s\-_./:]+/);
  for (const t of tokens) {
    let s = 0;
    if (nameWords.some((w) => w.startsWith(t))) s = 30;
    else if (name.includes(t)) s = 18;
    else if (id.includes(t)) s = 15;
    else if (lbl.includes(t)) s = 25;
    else if (family.includes(t)) s = 10;
    else if (vendor.includes(t)) s = 8;
    else if (prov.includes(t)) s = 8;
    else if (t.length >= 4 && (fuzzy(t, name.replace(/\s+/g, "")) || fuzzy(t, id))) s = 4;
    if (s === 0) return -1;
    score += s;
  }
  return score;
}

export function searchModels(models: PolyModel[], query: string | SearchIntent, sctx: SearchContext = {}): SearchResult[] {
  const intent = typeof query === "string" ? parseSearchQuery(query) : query;
  const now = sctx.now ?? sctx.ctx?.now ?? Date.now();
  const ctx = sctx.ctx ?? buildBadgeContext(models, now);
  const recents = sctx.recents instanceof Set ? sctx.recents : new Set(sctx.recents ?? []);
  const whole = intent.text.join(" ");
  const blendedMax = Math.max(1, ...models.map((m) => blendedPrice(m) ?? 0));
  const ctxMax = Math.max(1, ...models.map((m) => m.limits?.contextTokens ?? 0));

  const out: SearchResult[] = [];
  for (const m of models) {
    const matches: string[] = [];
    // Hard filters ----------------------------------------------------------
    if (intent.brands.length) {
      const hay = `${m.id} ${m.displayName} ${m.family ?? ""} ${m.metadata?.vendor ?? ""}`;
      const ok = intent.brands.some((b) => (b.provider && m.provider === b.provider) || b.match.test(hay));
      if (!ok) continue;
      matches.push(intent.brands.map((b) => b.label).join("/"));
    }
    if (intent.capabilities.some((c) => !m.capabilities[c])) continue;
    for (const c of intent.capabilities) matches.push(CAP_WORDS.find((w) => w.cap === c)?.chip ?? c);
    if (intent.minContext !== undefined) {
      if ((m.limits?.contextTokens ?? 0) < intent.minContext) continue;
      matches.push(`${fmtCtx(m.limits!.contextTokens!)} context`);
    }
    if (intent.maxInputPrice !== undefined) {
      const p = m.pricing?.inputPerMillion;
      if (typeof p !== "number" || p > intent.maxInputPrice) continue;
      matches.push(`$${p}/M in`);
    }
    if (intent.maxOutputPrice !== undefined) {
      const p = m.pricing?.outputPerMillion;
      if (typeof p !== "number" || p > intent.maxOutputPrice) continue;
      matches.push(`$${p}/M out`);
    }
    if (intent.openWeights) {
      if (!isOpenWeightsModel(m)) continue;
      matches.push("open weights");
    }
    if (intent.newest && !isNewModel(m, ctx) && intent.text.length === 0 && !intent.brands.length) continue;

    // Text --------------------------------------------------------------------
    const ts = textScore(m, intent.text, whole, sctx.labels?.[m.key], sctx.providerNames?.[m.provider] ?? m.provider);
    if (ts < 0) continue;
    let score = ts;

    // Soft boosts ------------------------------------------------------------
    const blended = blendedPrice(m);
    const fast = isFastModel(m, ctx);
    const frontier = isFrontierModel(m);
    if (intent.sort === "cheapest") {
      score += blended === null ? -10 : (1 - blended / blendedMax) * 30;
      if (blended !== null && blended <= ctx.cheapThreshold) matches.push("cheap");
    }
    if (intent.sort === "fastest") {
      score += fast ? 25 : 0;
      if (m.provider === "cerebras") score += 8;
      if (fast) matches.push("fast tier");
    }
    if (intent.sort === "context") score += ((m.limits?.contextTokens ?? 0) / ctxMax) * 15;
    if (intent.sort === "newest") {
      const t = firstSeenMs(m) ?? releaseDateMs(m);
      if (t !== null) score += Math.max(0, 20 - (now - t) / 86_400_000 / 10);
    }
    if (intent.quality) {
      if (frontier) {
        score += 25;
        matches.push("frontier");
      }
      if (m.capabilities.reasoning) score += 5;
      if (fast) score -= 8;
    }
    if (intent.coding) {
      if (isCodingModel(m)) {
        score += 30;
        matches.push("coding");
      } else if (frontier && m.capabilities.tools) score += 10;
      else if (!m.capabilities.tools) score -= 10;
    }
    // Base hygiene
    score += Math.min(10, Math.max(0, sortWeightOf(m) / 10));
    if (sctx.favorites?.has(m.key)) score += 6;
    if (recents.has(m.key)) score += 3;
    if (m.status === "deprecated") score -= 50;
    else if (m.status === "preview") score -= 2;
    else if (m.status === "unknown") score -= 4;
    if (m.provider === "openrouter" && intent.text.length) score -= 5; // prefer native providers for the same model

    out.push({ model: m, score, matches: [...new Set(matches)] });
  }
  out.sort((a, b) => b.score - a.score || sortWeightOf(b.model) - sortWeightOf(a.model) || a.model.displayName.localeCompare(b.model.displayName));
  return out;
}