import type { PolyModel } from "@/lib/ai/core/types";
import { analyzePrompt, routeModels, type RouterMode } from "@/lib/client/router";
import { blendedPrice, firstSeenMs, isCodingModel, isFastModel, isFrontierModel, isNewModel, isOpenWeightsModel, releaseDateMs, sortWeightOf, type BadgeContext } from "./badges";

/** Sections of the model picker (in display order). */
export type SectionKey = "favorites" | "recent" | "recommended" | "fast" | "reasoning" | "cheapest" | "context" | "vision" | "coding" | "open-weights" | "new";

export interface SectionMeta {
  key: SectionKey;
  label: string;
  hint: string;
}

export const SECTION_META: SectionMeta[] = [
  { key: "favorites", label: "Favorites", hint: "Starred by you" },
  { key: "recent", label: "Recently used", hint: "Your last picks" },
  { key: "recommended", label: "Recommended", hint: "Smart Router picks for this prompt" },
  { key: "fast", label: "Fast", hint: "Low-latency tiers" },
  { key: "reasoning", label: "Best reasoning", hint: "Extended thinking" },
  { key: "cheapest", label: "Cheapest", hint: "Lowest blended price" },
  { key: "context", label: "Largest context", hint: "Biggest windows" },
  { key: "vision", label: "Vision", hint: "Understand images" },
  { key: "coding", label: "Coding", hint: "Tuned for code" },
  { key: "open-weights", label: "Open source", hint: "Open-weights models" },
  { key: "new", label: "New models", hint: "Added in the last 30 days" },
];

export interface Section extends SectionMeta {
  /** First `perSection` items. */
  items: PolyModel[];
  /** Every match (for "Show all"). */
  all: PolyModel[];
  /** Total before `perSection` truncation. */
  total: number;
}

export interface SectionOptions {
  favorites: Set<string>;
  recents: string[];
  ctx: BadgeContext;
  mode?: RouterMode;
  /** Current composer draft — drives "Recommended". */
  draft?: string;
  perSection?: number;
}

const GENERIC_PROMPT = "Help me think through a question and write a clear, well-structured answer.";

const byWeight = (a: PolyModel, b: PolyModel) => sortWeightOf(b) - sortWeightOf(a) || a.displayName.localeCompare(b.displayName);

export function recommendedModels(models: PolyModel[], opts: { favorites?: Set<string>; mode?: RouterMode; draft?: string; limit?: number }) {
  const analysis = analyzePrompt(opts.draft?.trim() || GENERIC_PROMPT);
  const route = routeModels(models, analysis, opts.mode ?? "balanced", { favorites: opts.favorites });
  const list = [route.recommended, ...route.alternatives].filter((c): c is NonNullable<typeof c> => Boolean(c));
  return { route, models: list.slice(0, opts.limit ?? 3).map((c) => c.model) };
}

/** Build every non-empty section. `models` must already be filtered to what the user can pick. */
export function buildSelectorSections(models: PolyModel[], opts: SectionOptions): Section[] {
  const per = opts.perSection ?? 5;
  const live = models.filter((m) => m.status !== "deprecated");
  const byKey = new Map(models.map((m) => [m.key, m]));

  const lists: Record<SectionKey, PolyModel[]> = {
    favorites: models.filter((m) => opts.favorites.has(m.key)).sort(byWeight),
    recent: opts.recents.map((k) => byKey.get(k)).filter((m): m is PolyModel => Boolean(m)),
    recommended: recommendedModels(live, { favorites: opts.favorites, mode: opts.mode, draft: opts.draft, limit: 3 }).models,
    fast: live.filter((m) => isFastModel(m, opts.ctx)).sort((a, b) => (a.pricing?.outputPerMillion ?? 99) - (b.pricing?.outputPerMillion ?? 99) || byWeight(a, b)),
    reasoning: live.filter((m) => m.capabilities.reasoning).sort((a, b) => Number(isFrontierModel(b)) - Number(isFrontierModel(a)) || byWeight(a, b)),
    cheapest: live.filter((m) => blendedPrice(m) !== null).sort((a, b) => (blendedPrice(a) ?? 0) - (blendedPrice(b) ?? 0) || byWeight(a, b)),
    context: live.filter((m) => (m.limits?.contextTokens ?? 0) > 0).sort((a, b) => (b.limits?.contextTokens ?? 0) - (a.limits?.contextTokens ?? 0) || byWeight(a, b)),
    vision: live.filter((m) => m.capabilities.vision).sort(byWeight),
    coding: live.filter(isCodingModel).sort(byWeight),
    "open-weights": live.filter(isOpenWeightsModel).sort(byWeight),
    new: live.filter((m) => isNewModel(m, opts.ctx)).sort((a, b) => (firstSeenMs(b) ?? releaseDateMs(b) ?? 0) - (firstSeenMs(a) ?? releaseDateMs(a) ?? 0)),
  };

  return SECTION_META.map((meta) => {
    const all = lists[meta.key];
    return { ...meta, items: all.slice(0, per), all, total: all.length };
  }).filter((s) => s.total > 0);
}