import { describe, it, expect } from "vitest";
import type { PolyModel } from "@/lib/ai/core/types";
import { parseSearchQuery, searchModels, isEmptyIntent } from "@/lib/models/search";
import { buildBadgeContext } from "@/lib/models/badges";

const caps = (over: Partial<PolyModel["capabilities"]> = {}): PolyModel["capabilities"] => ({ text: true, vision: false, audioInput: false, audioOutput: false, imageGeneration: false, video: false, reasoning: false, tools: true, structuredOutput: false, streaming: true, files: false, webSearch: false, ...over });

const NOW = Date.parse("2026-09-11T12:00:00Z");
const daysAgo = (d: number) => new Date(NOW - d * 86_400_000).toISOString();

function model(p: Partial<PolyModel> & { id: string; provider: PolyModel["provider"] }): PolyModel {
  return { key: `${p.provider}/${p.id}`, displayName: p.id, capabilities: caps(), parameters: {}, status: "active", pricing: null, metadata: { sortWeight: 0, firstSeenAt: daysAgo(200) }, ...p };
}

const MODELS: PolyModel[] = [
  model({ id: "gpt-5.5", provider: "openai", displayName: "GPT-5.5", family: "GPT-5.5", capabilities: caps({ vision: true, reasoning: true, structuredOutput: true, files: true, webSearch: true }), limits: { contextTokens: 1_050_000, maxOutputTokens: 128_000 }, pricing: { inputPerMillion: 5, outputPerMillion: 30 }, metadata: { sortWeight: 90, firstSeenAt: daysAgo(200) } }),
  model({ id: "gpt-5.4-mini", provider: "openai", displayName: "GPT-5.4 mini", family: "GPT-5.4", capabilities: caps({ vision: true, reasoning: true, structuredOutput: true }), limits: { contextTokens: 400_000 }, pricing: { inputPerMillion: 0.4, outputPerMillion: 1.6 }, metadata: { sortWeight: 82, firstSeenAt: daysAgo(200) } }),
  model({ id: "claude-sonnet-5", provider: "anthropic", displayName: "Claude Sonnet 5", family: "Sonnet", capabilities: caps({ vision: true, reasoning: true, files: true, webSearch: true }), limits: { contextTokens: 1_000_000, maxOutputTokens: 64_000 }, pricing: { inputPerMillion: 2, outputPerMillion: 10 }, metadata: { sortWeight: 90, firstSeenAt: daysAgo(10) } }),
  model({ id: "claude-haiku-4-5", provider: "anthropic", displayName: "Claude Haiku 4.5", family: "Haiku", capabilities: caps({ vision: true }), limits: { contextTokens: 200_000 }, pricing: { inputPerMillion: 1, outputPerMillion: 5 }, metadata: { sortWeight: 60, firstSeenAt: daysAgo(200) } }),
  model({ id: "gemini-3.8-flash", provider: "gemini", displayName: "Gemini 3.8 Flash", family: "Gemini", capabilities: caps({ vision: true, reasoning: true, structuredOutput: true, webSearch: true }), limits: { contextTokens: 1_048_576 }, pricing: { inputPerMillion: 0.3, outputPerMillion: 2.5 }, metadata: { sortWeight: 100, firstSeenAt: daysAgo(200) } }),
  model({ id: "gemini-3.1-pro", provider: "gemini", displayName: "Gemini 3.1 Pro", family: "Gemini", capabilities: caps({ vision: true, reasoning: true, structuredOutput: true }), limits: { contextTokens: 2_000_000 }, pricing: { inputPerMillion: 2, outputPerMillion: 12 }, metadata: { sortWeight: 96, firstSeenAt: daysAgo(200) } }),
  model({ id: "codestral-latest", provider: "mistral", displayName: "Codestral", family: "Codestral", capabilities: caps({ structuredOutput: true }), limits: { contextTokens: 256_000 }, pricing: { inputPerMillion: 0.3, outputPerMillion: 0.9 }, metadata: { sortWeight: 50, firstSeenAt: daysAgo(200) } }),
  model({ id: "meta-llama/llama-3.1-8b-instruct", provider: "openrouter", displayName: "Meta: Llama 3.1 8B Instruct", family: "Meta", capabilities: caps(), limits: { contextTokens: 131_072 }, pricing: { inputPerMillion: 0.02, outputPerMillion: 0.05 }, metadata: { sortWeight: 20, vendor: "meta-llama", firstSeenAt: daysAgo(200) } }),
  model({ id: "deepseek-v4-flash", provider: "deepseek", displayName: "DeepSeek V4 Flash", family: "DeepSeek", capabilities: caps({ reasoning: true }), limits: { contextTokens: 128_000 }, pricing: { inputPerMillion: 0.14, outputPerMillion: 0.28 }, metadata: { sortWeight: 70, firstSeenAt: daysAgo(200) } }),
  model({ id: "old-model", provider: "openai", displayName: "Old Model", status: "deprecated", pricing: { inputPerMillion: 1, outputPerMillion: 2 }, capabilities: caps({ vision: true }) }),
];

const ctx = buildBadgeContext(MODELS, NOW);
const ids = (r: { model: PolyModel }[]) => r.map((x) => x.model.id);

describe("parseSearchQuery", () => {
  it("understands capability + cost intents", () => {
    const i = parseSearchQuery("cheap vision model");
    expect(i.capabilities).toEqual(["vision"]);
    expect(i.sort).toBe("cheapest");
    expect(i.text).toEqual([]);
    expect(i.chips).toContain("Vision");
  });
  it("parses context sizes", () => {
    expect(parseSearchQuery("1M context").minContext).toBe(1_000_000);
    expect(parseSearchQuery("200k").minContext).toBe(200_000);
    expect(parseSearchQuery("128k ctx").minContext).toBe(128_000);
    expect(parseSearchQuery("long context").minContext).toBe(400_000);
  });
  it("parses price caps", () => {
    expect(parseSearchQuery("under $1/M").maxInputPrice).toBe(1);
    expect(parseSearchQuery("< $2 per million output").maxOutputPrice).toBe(2);
    expect(parseSearchQuery("0.5$").maxInputPrice).toBe(0.5);
  });
  it("recognises brands, json schema, open source and reasoning", () => {
    const g = parseSearchQuery("fastest gemini");
    expect(g.brands.map((b) => b.key)).toEqual(["gemini"]);
    expect(g.sort).toBe("fastest");
    expect(parseSearchQuery("json schema").capabilities).toEqual(["structuredOutput"]);
    expect(parseSearchQuery("open source").openWeights).toBe(true);
    expect(parseSearchQuery("open-weights coder").coding).toBe(true);
    expect(parseSearchQuery("reasoning").capabilities).toEqual(["reasoning"]);
    expect(parseSearchQuery("OpenAI").brands[0].provider).toBe("openai");
  });
  it("keeps free text and detects empty intents", () => {
    expect(parseSearchQuery("sonnet").text).toEqual(["sonnet"]);
    expect(isEmptyIntent(parseSearchQuery("   "))).toBe(true);
    expect(isEmptyIntent(parseSearchQuery("the model"))).toBe(true);
  });
});

describe("searchModels", () => {
  it("filters by capability and ranks cheapest first", () => {
    const r = searchModels(MODELS, "cheap vision model", { ctx });
    expect(r.every((x) => x.model.capabilities.vision)).toBe(true);
    expect(r[0].model.id).toBe("gemini-3.8-flash");
    expect(ids(r)).not.toContain("codestral-latest");
    // deprecated models sink to the bottom
    expect(r[r.length - 1].model.id).toBe("old-model");
  });
  it("applies context and price constraints", () => {
    expect(ids(searchModels(MODELS, "1M context", { ctx }))).toEqual(expect.arrayContaining(["gpt-5.5", "claude-sonnet-5", "gemini-3.8-flash", "gemini-3.1-pro"]));
    expect(ids(searchModels(MODELS, "1M context", { ctx }))).not.toContain("gpt-5.4-mini");
    const cheap = searchModels(MODELS, "under $1/M", { ctx });
    expect(cheap.every((x) => (x.model.pricing?.inputPerMillion ?? 99) <= 1)).toBe(true);
  });
  it("matches brands across native providers and gateways", () => {
    const g = searchModels(MODELS, "fastest gemini", { ctx });
    expect(ids(g)).toEqual(["gemini-3.8-flash", "gemini-3.1-pro"]);
    const llama = searchModels(MODELS, "llama", { ctx });
    expect(ids(llama)).toEqual(["meta-llama/llama-3.1-8b-instruct"]);
  });
  it("ranks exact and prefix name matches first and tolerates a typo", () => {
    expect(searchModels(MODELS, "GPT-5.5", { ctx })[0].model.id).toBe("gpt-5.5");
    expect(searchModels(MODELS, "sonnet", { ctx })[0].model.id).toBe("claude-sonnet-5");
    expect(searchModels(MODELS, "sonet", { ctx })[0].model.id).toBe("claude-sonnet-5");
    expect(searchModels(MODELS, "zzzz", { ctx })).toHaveLength(0);
  });
  it("uses user labels and favorites", () => {
    const r = searchModels(MODELS, "research", { ctx, labels: { "anthropic/claude-sonnet-5": "My research model" } });
    expect(ids(r)).toEqual(["claude-sonnet-5"]);
    const fav = searchModels(MODELS, "gemini", { ctx, favorites: new Set(["gemini/gemini-3.1-pro"]) });
    expect(fav[0].model.id).toBe("gemini-3.1-pro");
  });
  it("handles open source, coding, json schema and new", () => {
    expect(ids(searchModels(MODELS, "open source", { ctx }))).toEqual(expect.arrayContaining(["meta-llama/llama-3.1-8b-instruct", "deepseek-v4-flash"]));
    expect(searchModels(MODELS, "coding", { ctx })[0].model.id).toBe("codestral-latest");
    expect(searchModels(MODELS, "json schema", { ctx }).every((x) => x.model.capabilities.structuredOutput)).toBe(true);
    expect(ids(searchModels(MODELS, "new", { ctx }))).toEqual(["claude-sonnet-5"]);
  });
});