import type { PolyModel, ProviderId } from "@/lib/client/types";
import { estimateTextTokens, estimateAttachmentTokens, estimateCost, type AttachmentLike } from "./tokens";

/**
 * PolyLLM Smart Router — transparent, client-side model recommendation.
 *
 * Pure function over the user's usable models: no network, no hidden state. The UI always shows
 * the recommendation, the reasons and the estimated cost before anything is sent, and never
 * silently upgrades to an expensive model (see `budgetCap`).
 */
export type RouterMode = "balanced" | "quality" | "fastest" | "cheapest" | "coding" | "research";

export const ROUTER_MODES: { value: RouterMode; label: string; description: string }[] = [
  { value: "balanced", label: "Balanced", description: "Good quality at a sensible price" },
  { value: "quality", label: "Best quality", description: "Strongest reasoning, cost secondary" },
  { value: "fastest", label: "Fastest", description: "Lowest latency tiers" },
  { value: "cheapest", label: "Lowest cost", description: "Cheapest model that fits the task" },
  { value: "coding", label: "Coding", description: "Code generation, review and debugging" },
  { value: "research", label: "Research", description: "Long context, reasoning, web search" },
];

export type TaskType = "coding" | "research" | "writing" | "reasoning" | "analysis" | "vision" | "translation" | "chat";

export interface PromptAnalysis {
  task: TaskType;
  inputTokens: number;
  needsVision: boolean;
  needsFiles: boolean;
  needsLongContext: boolean;
  needsReasoning: boolean;
  needsWeb: boolean;
  needsJson: boolean;
  signals: string[];
}

const CODE_RE = /```|\bfunction\b|\bconst\b|\bimport\b|\bdef\b|\bclass\b|\breturn\b|=>|\bSELECT\b|\bnpm\b|\bpip\b|stack ?trace|exception|bug|refactor|typescript|python|rust|\bsql\b|regex|compile|debug/i;
const RESEARCH_RE = /research|sources?|cite|citation|literature|paper|study|studies|latest|news|today|current|2026|compare .* (and|vs)|pros and cons|state of the art|survey/i;
const REASON_RE = /prove|proof|derive|step[- ]by[- ]step|theorem|puzzle|riddle|optimi[sz]e|strategy|plan (a|the)|trade-?offs?|why does|analy[sz]e deeply|chain of thought|logic/i;
const WRITING_RE = /write|draft|essay|blog|email|letter|story|poem|rewrite|paraphrase|tone|copy(writing)?|headline|slogan|summar/i;
const ANALYSIS_RE = /csv|spreadsheet|dataset|data\b|table|statistics|trend|chart|forecast|financial|revenue|kpi|metrics/i;
const TRANSLATE_RE = /translate|traduis|traduction|in (french|spanish|german|italian|japanese|chinese|english)/i;
const JSON_RE = /\bjson\b|schema|structured output|as a list of objects|key[- ]value/i;
const WEB_RE = /search the web|look up|latest|news|today|this week|current price|weather|stock|who won|release date/i;

export function analyzePrompt(text: string, attachments: AttachmentLike[] = [], opts: { historyTokens?: number; systemPrompt?: string | null; webSearch?: boolean; responseFormat?: boolean } = {}): PromptAnalysis {
  const signals: string[] = [];
  const t = text ?? "";
  const attTokens = attachments.reduce((n, a) => n + estimateAttachmentTokens(a), 0);
  const inputTokens = estimateTextTokens(t) + estimateTextTokens(opts.systemPrompt ?? "") + attTokens + (opts.historyTokens ?? 0);
  const needsVision = attachments.some((a) => a.kind === "image");
  const needsFiles = attachments.some((a) => a.kind === "pdf");
  const needsLongContext = inputTokens > 60_000;
  const needsWeb = Boolean(opts.webSearch) || WEB_RE.test(t);
  const needsJson = Boolean(opts.responseFormat) || JSON_RE.test(t);
  if (needsVision) signals.push("image attached");
  if (needsFiles) signals.push("PDF attached");
  if (needsLongContext) signals.push(`${Math.round(inputTokens / 1000)}K tokens of context`);
  if (needsWeb) signals.push("fresh information");
  if (needsJson) signals.push("structured output");

  let task: TaskType = "chat";
  const codeScore = (t.match(CODE_RE) ?? []).length + (t.includes("```") ? 3 : 0);
  if (attachments.some((a) => a.kind === "code")) signals.push("code file");
  if (codeScore >= 2 || attachments.some((a) => a.kind === "code")) task = "coding";
  else if (TRANSLATE_RE.test(t)) task = "translation";
  else if (REASON_RE.test(t)) task = "reasoning";
  else if (RESEARCH_RE.test(t)) task = "research";
  else if (ANALYSIS_RE.test(t) || attachments.some((a) => a.kind === "csv" || a.kind === "json")) task = "analysis";
  else if (WRITING_RE.test(t)) task = "writing";
  else if (needsVision) task = "vision";
  const needsReasoning = task === "reasoning" || task === "coding" || task === "analysis" || (task === "research" && t.length > 400) || t.length > 1500;
  if (task !== "chat") signals.unshift(`${task} task`);
  if (needsReasoning && task !== "reasoning") signals.push("benefits from reasoning");
  return { task, inputTokens, needsVision, needsFiles, needsLongContext, needsReasoning, needsWeb, needsJson, signals };
}

export interface RouteCandidate {
  model: PolyModel;
  score: number;
  estimatedUsd: number | null;
  reasons: string[];
}

export interface RouteResult {
  analysis: PromptAnalysis;
  recommended: RouteCandidate | null;
  alternatives: RouteCandidate[];
  mode: RouterMode;
  /** True when the recommendation exceeds the cap and needs explicit confirmation. */
  needsConfirmation: boolean;
}

const FAST_RE = /flash|mini|nano|haiku|fast|lite|non-reasoning|instant|turbo|small|ministral|8b|9b|gemma|gpt-oss-20b|llama-3\.1-8b/i;
const FRONTIER_RE = /opus|fable|gpt-5\.5|gpt-6|pro\b|grok-4(\.\d+)?(?!-fast)|large|k3|deepseek-v4-pro|magistral-medium/i;
const CODING_RE = /codestral|code|coder|devstral|gpt-5\.5|sonnet|opus|fable|k2\.7|deepseek/i;
const OPEN_RE = /llama|qwen|gemma|gpt-oss|mistral-small|mixtral|deepseek|kimi|k2|k3|glm|command/i;

const PROVIDER_LATENCY_BONUS: Partial<Record<ProviderId, number>> = { cerebras: 3, xai: 0.5, gemini: 0.5, openrouter: -0.3 };

/** Median USD per 1M output tokens across usable models — used as the "expensive" reference. */
function medianOut(models: PolyModel[]): number {
  const v = models.map((m) => m.pricing?.outputPerMillion).filter((x): x is number => typeof x === "number").sort((a, b) => a - b);
  return v.length ? v[Math.floor(v.length / 2)] : 5;
}

export function routeModels(models: PolyModel[], analysis: PromptAnalysis, mode: RouterMode = "balanced", opts: { budgetCapUsd?: number; favorites?: Set<string>; expectedOutputTokens?: number } = {}): RouteResult {
  const expectedOut = opts.expectedOutputTokens ?? (analysis.task === "writing" || analysis.task === "coding" ? 1200 : 600);
  const usable = models.filter((m) => m.status !== "deprecated" && m.capabilities.text !== false);
  const median = medianOut(usable);
  const cap = opts.budgetCapUsd ?? 0.5;

  const scored: RouteCandidate[] = usable
    .filter((m) => {
      if (analysis.needsVision && !m.capabilities.vision) return false;
      if (analysis.needsFiles && !m.capabilities.files) return false;
      const ctx = m.limits?.contextTokens ?? 128_000;
      if (analysis.inputTokens + expectedOut > ctx * 0.92) return false;
      return true;
    })
    .map((m) => {
      const reasons: string[] = [];
      let score = 0;
      const id = `${m.id} ${m.displayName}`;
      const fast = FAST_RE.test(id);
      const frontier = FRONTIER_RE.test(id);
      const outPrice = m.pricing?.outputPerMillion ?? median;
      const inPrice = m.pricing?.inputPerMillion ?? median / 4;
      const cost = estimateCost(m, analysis.inputTokens, expectedOut).usd;
      const cheapness = Math.max(0, 1 - outPrice / (median * 2)); // 0..1
      const ctx = m.limits?.contextTokens ?? 0;

      // Capability fit
      if (analysis.needsReasoning && m.capabilities.reasoning) {
        score += 3;
        reasons.push("reasoning");
      } else if (analysis.needsReasoning) score -= 2;
      if (analysis.needsWeb && m.capabilities.webSearch) {
        score += 2.5;
        reasons.push("web search");
      } else if (analysis.needsWeb) score -= 1;
      if (analysis.needsJson && m.capabilities.structuredOutput) {
        score += 1.5;
        reasons.push("JSON schema");
      }
      if (analysis.needsVision) reasons.push("vision");
      if (analysis.needsLongContext && ctx >= 400_000) {
        score += 2.5;
        reasons.push(`${Math.round(ctx / 1000)}K context`);
      } else if (ctx >= 1_000_000 && (mode === "research" || analysis.task === "research")) {
        score += 1;
        reasons.push("1M context");
      }
      if ((analysis.task === "coding" || mode === "coding") && CODING_RE.test(id)) {
        score += 2;
        reasons.push("strong at code");
      }
      if (analysis.task === "translation" && fast) score += 1;
      if (analysis.task === "chat" && fast) score += 1.5;

      // Tier
      if (frontier) {
        if (mode === "quality" || analysis.task === "reasoning" || analysis.task === "research") {
          score += 3;
          reasons.push("frontier quality");
        } else if (mode === "balanced" && analysis.needsReasoning) score += 1;
        else if (mode === "cheapest" || mode === "fastest") score -= 3;
      }
      if (fast) {
        if (mode === "fastest") {
          score += 3.5;
          reasons.push("fast tier");
        } else if (mode === "cheapest") score += 1.5;
        else if (mode === "quality") score -= 2;
        else if (mode === "balanced" && !analysis.needsReasoning) {
          score += 1;
          reasons.push("fast");
        }
      }
      score += (PROVIDER_LATENCY_BONUS[m.provider] ?? 0) * (mode === "fastest" ? 1 : 0.2);

      // Cost
      const costWeight = mode === "cheapest" ? 5 : mode === "balanced" ? 2.2 : mode === "quality" ? 0.4 : mode === "fastest" ? 1 : 1.5;
      score += cheapness * costWeight;
      if (cheapness > 0.7) reasons.push(`cheap (${inPrice.toFixed(2)}/${outPrice.toFixed(2)} per 1M)`);
      if (cost !== null && cost > cap && mode !== "quality") score -= 2;

      // Preference & hygiene
      if (opts.favorites?.has(m.key)) {
        score += 0.8;
        reasons.push("favorite");
      }
      if (m.status === "preview") score -= 0.3;
      if (m.provider === "openrouter") score -= 0.6; // prefer native connections when the same model exists
      if (OPEN_RE.test(id) && mode === "quality") score -= 0.5;

      return { model: m, score, estimatedUsd: cost, reasons: dedupe(reasons).slice(0, 3) };
    })
    .sort((a, b) => b.score - a.score);

  const recommended = scored[0] ?? null;
  const alternatives = scored.slice(1, 4);
  const needsConfirmation = Boolean(recommended?.estimatedUsd && recommended.estimatedUsd > cap);
  return { analysis, recommended, alternatives, mode, needsConfirmation };
}

function dedupe<T>(arr: T[]): T[] {
  return [...new Set(arr)];
}

/** Human explanation: "Strong reasoning + 1M context". */
export function explainRoute(c: RouteCandidate | null): string {
  if (!c) return "No compatible model is connected.";
  if (!c.reasons.length) return "Best overall fit for this prompt.";
  return c.reasons.map((r, i) => (i === 0 ? r.charAt(0).toUpperCase() + r.slice(1) : r)).join(" + ");
}