/**
 * Benchmark-driven model selection.
 *
 * When `models.<unit>` is not set in preferences, this module picks the
 * best-scoring model from the allow-listed providers for each unit type.
 * Scoring is a weighted combination of published benchmarks
 * (`learning/data/model-benchmarks.json`) with per-unit-type profiles
 * that emphasise the dimensions that actually matter for that work:
 *   - plan-milestone / plan-slice    → reasoning-heavy (hle, aime, gpqa)
 *   - research-*                      → mixed (mmlu_pro, browse_comp, ...)
 *   - execute-task (heavy)            → coding (swe_bench, live_code_bench)
 *   - execute-task (light/standard)   → coding + instruction following
 *   - complete-* / execution_simple   → fast+correct (human_eval, ifeval)
 *   - gate-evaluate / validate-*      → reasoning + coding
 *
 * Missing benchmark scores are treated as 0 (model ranked last rather
 * than excluded) so freshly-launched models without benchmark data are
 * still dispatchable — they just don't displace an already-ranked peer.
 *
 * This is the inner primitive behind the "auto-benchmark" preference mode
 * users select by leaving `models.*` empty.
 */

import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import { dirname } from "node:path";

// ─── Types ───────────────────────────────────────────────────────────────────

type BenchmarkKey =
  | "swe_bench" | "swe_bench_verified" | "live_code_bench" | "human_eval"
  | "hle" | "aime_2026" | "gpqa" | "mmlu_pro" | "bbh"
  | "browse_comp" | "simple_qa" | "long_context_ruler"
  | "arena_elo" | "instruction_following";

interface BenchmarkRecord {
  [k: string]: number | string | null | undefined;
  context_window?: number;
  max_output_tokens?: number;
}

interface BenchmarkData {
  _meta?: unknown;
  [modelKey: string]: BenchmarkRecord | unknown;
}

export interface CandidateModel {
  /** Provider ID (e.g. "kimi-coding", "mistral", "opencode-go") */
  provider: string;
  /** Bare model ID without provider prefix (e.g. "k2p5", "codestral-latest") */
  id: string;
}

export interface BenchmarkSelectionResult {
  primary: string;               // "provider/model-id"
  fallbacks: string[];           // ordered, deduplicated
  /** Raw per-model score, exposed for logging / UI */
  scores: Record<string, number>;
  /** Number of benchmark dimensions actually populated for the top model */
  topCoverage: number;
  /** Unit-type label (for debug logs) */
  profile: string;
}

// ─── Benchmark File Loader ───────────────────────────────────────────────────

let _benchmarksCache: BenchmarkData | null = null;

function loadBenchmarks(): BenchmarkData {
  if (_benchmarksCache) return _benchmarksCache;
  const __filename = fileURLToPath(import.meta.url);
  const here = dirname(__filename);
  // Works for both .ts (dev) and .js (dist) since we copy the data file 1:1.
  const path = join(here, "learning", "data", "model-benchmarks.json");
  if (!existsSync(path)) {
    _benchmarksCache = {};
    return _benchmarksCache;
  }
  try {
    _benchmarksCache = JSON.parse(readFileSync(path, "utf-8")) as BenchmarkData;
  } catch {
    _benchmarksCache = {};
  }
  return _benchmarksCache;
}

/** Testing: reset the in-memory benchmark cache. */
export function _resetBenchmarkCache(): void {
  _benchmarksCache = null;
}

// ─── Unit-Type → Weight Profiles ─────────────────────────────────────────────
//
// Weights sum to ~1.0 per profile; small deviations are tolerated (we
// normalise by the total of populated weights so a model missing some
// dimensions isn't artificially dragged down vs a model missing others).

type WeightProfile = Partial<Record<BenchmarkKey, number>>;

const PROFILES: Record<string, { weights: WeightProfile; label: string }> = {
  // Planning in SF is agent-style decomposition work, not pure math
  // olympiad reasoning. Weight swe_bench (agent/coding reasoning) and
  // live_code_bench heavier; keep hle/gpqa for general capability.
  "plan-milestone":      { weights: { swe_bench: 0.25, live_code_bench: 0.20, hle: 0.15, gpqa: 0.15, mmlu_pro: 0.15, aime_2026: 0.10 }, label: "agent-planning" },
  "plan-slice":          { weights: { swe_bench: 0.25, live_code_bench: 0.20, hle: 0.15, gpqa: 0.15, mmlu_pro: 0.15, aime_2026: 0.10 }, label: "agent-planning" },
  "replan-slice":        { weights: { hle: 0.25, gpqa: 0.20, swe_bench: 0.30, mmlu_pro: 0.15, instruction_following: 0.10 }, label: "replanning" },
  "discuss":             { weights: { hle: 0.25, mmlu_pro: 0.25, gpqa: 0.20, instruction_following: 0.15, simple_qa: 0.15 }, label: "discussion" },
  "discuss-milestone":   { weights: { hle: 0.25, mmlu_pro: 0.25, gpqa: 0.20, instruction_following: 0.15, simple_qa: 0.15 }, label: "discussion" },
  "discuss-slice":       { weights: { hle: 0.25, mmlu_pro: 0.25, gpqa: 0.20, instruction_following: 0.15, simple_qa: 0.15 }, label: "discussion" },
  "discuss-headless":    { weights: { hle: 0.25, mmlu_pro: 0.25, gpqa: 0.20, instruction_following: 0.15, simple_qa: 0.15 }, label: "discussion" },
  "research-milestone":  { weights: { mmlu_pro: 0.25, hle: 0.20, human_eval: 0.20, browse_comp: 0.15, simple_qa: 0.10, gpqa: 0.10 }, label: "research" },
  "research-slice":      { weights: { mmlu_pro: 0.25, hle: 0.20, human_eval: 0.20, browse_comp: 0.15, simple_qa: 0.10, gpqa: 0.10 }, label: "research" },
  "execute-task":        { weights: { swe_bench: 0.35, swe_bench_verified: 0.25, live_code_bench: 0.20, human_eval: 0.15, instruction_following: 0.05 }, label: "coding" },
  "reactive-execute":    { weights: { swe_bench: 0.30, live_code_bench: 0.25, human_eval: 0.20, hle: 0.15, instruction_following: 0.10 }, label: "coding" },
  "execute-task-simple": { weights: { human_eval: 0.40, instruction_following: 0.35, long_context_ruler: 0.25 }, label: "fast+correct" },
  "execution_simple":    { weights: { human_eval: 0.40, instruction_following: 0.35, long_context_ruler: 0.25 }, label: "fast+correct" },
  "complete-slice":      { weights: { instruction_following: 0.40, human_eval: 0.35, long_context_ruler: 0.25 }, label: "fast+correct" },
  "complete-milestone":  { weights: { instruction_following: 0.40, human_eval: 0.35, long_context_ruler: 0.25 }, label: "fast+correct" },
  "gate-evaluate":       { weights: { swe_bench: 0.30, hle: 0.25, gpqa: 0.25, instruction_following: 0.20 }, label: "review" },
  "validate-milestone":  { weights: { hle: 0.30, gpqa: 0.25, mmlu_pro: 0.25, swe_bench: 0.20 }, label: "validation" },
  "subagent":            { weights: { swe_bench: 0.30, live_code_bench: 0.25, human_eval: 0.25, hle: 0.20 }, label: "subagent-default" },
  "run-uat":             { weights: { human_eval: 0.45, instruction_following: 0.40, long_context_ruler: 0.15 }, label: "uat" },
  "reassess-roadmap":    { weights: { mmlu_pro: 0.30, hle: 0.25, gpqa: 0.25, browse_comp: 0.10, simple_qa: 0.10 }, label: "reassessment" },
};

// Fallback for unit types not in the table — treat as standard coding.
const DEFAULT_PROFILE: WeightProfile = {
  swe_bench: 0.30, live_code_bench: 0.25, human_eval: 0.25, hle: 0.20,
};

function profileForUnitType(unitType: string): { weights: WeightProfile; label: string } {
  const direct = PROFILES[unitType];
  if (direct) return direct;
  // hook/* units inherit DEFAULT_PROFILE
  return { weights: DEFAULT_PROFILE, label: `default(${unitType})` };
}

// ─── Scoring ─────────────────────────────────────────────────────────────────

/**
 * Match a provider+model pair to a benchmark record key. Benchmarks are
 * keyed by bare model ID (e.g. "devstral-latest", "k2p5"), while registered
 * models may carry versioned suffixes (`devstral-2507`, `minimax-m2.7`).
 * We try exact match first, then strip common version/date suffixes, then
 * try a family-level key (e.g. `mistral-large-2411` → `mistral-large-latest`).
 */
function findBenchmarkKey(modelId: string, benchmarks: BenchmarkData): string | null {
  if (modelId in benchmarks) return modelId;
  // Strip date-style suffixes: "devstral-medium-2507" → "devstral-medium"
  const noDate = modelId.replace(/-\d{4}$/, "");
  if (noDate !== modelId) {
    if (noDate in benchmarks) return noDate;
    // Many vendors only publish benchmarks for the "-latest" alias.
    // "devstral-medium-2507" → "devstral-medium" → try "devstral-medium-latest".
    const latestAlias = `${noDate}-latest`;
    if (latestAlias in benchmarks) return latestAlias;
  }
  // Also try "-latest" alias when the model ID ends with a version number
  // (e.g. "minimax-m2.7" → look up "minimax-m2.7-latest", "minimax-m2-latest").
  const versionStripped = modelId.replace(/-\d+(\.\d+)?$/, "");
  if (versionStripped !== modelId) {
    const latestKey = `${versionStripped}-latest`;
    if (latestKey in benchmarks) return latestKey;
    if (versionStripped in benchmarks) return versionStripped;
  }
  // Case-insensitive match — last resort for casing drift (MiniMax-M2 vs
  // minimax-m2). The catalog uses one convention, the benchmark file
  // another.
  const lower = modelId.toLowerCase();
  for (const key of Object.keys(benchmarks)) {
    if (key === "_meta") continue;
    if (key.toLowerCase() === lower) return key;
  }
  return null;
}

// Some benchmarks are practical equivalents — vendors publish one or the
// other but rarely both. Treat them as fungible: whichever is populated
// fills the profile slot. This prevents MiniMax (publishes
// swe_bench_verified=80) from being penalised vs z.ai GLM-5.1 (publishes
// swe_bench=78) on a weight that references only "swe_bench".
const DIMENSION_EQUIVALENTS: Partial<Record<BenchmarkKey, BenchmarkKey[]>> = {
  swe_bench: ["swe_bench_verified"],
  swe_bench_verified: ["swe_bench"],
};

function readDimension(rec: BenchmarkRecord, dim: BenchmarkKey): number | null {
  const direct = rec[dim];
  if (typeof direct === "number" && Number.isFinite(direct)) return direct;
  const equivalents = DIMENSION_EQUIVALENTS[dim] ?? [];
  for (const alt of equivalents) {
    const v = rec[alt];
    if (typeof v === "number" && Number.isFinite(v)) return v;
  }
  return null;
}

function scoreCandidate(
  candidate: CandidateModel,
  profile: WeightProfile,
  benchmarks: BenchmarkData,
): { score: number; coverage: number } {
  const key = findBenchmarkKey(candidate.id, benchmarks);
  if (!key) return { score: 0, coverage: 0 };
  const rec = benchmarks[key] as BenchmarkRecord | undefined;
  if (!rec || typeof rec !== "object") return { score: 0, coverage: 0 };

  let weightedSum = 0;
  let weightTotal = 0;
  let profileTotal = 0;
  let coverage = 0;
  for (const [dim, weight] of Object.entries(profile)) {
    profileTotal += weight;
    const v = readDimension(rec, dim as BenchmarkKey);
    if (v !== null) {
      weightedSum += weight * v;
      weightTotal += weight;
      coverage++;
    }
  }
  if (weightTotal === 0) return { score: 0, coverage: 0 };
  // Normalise by populated weight so a model with 2 dimensions at 90 isn't
  // crushed by a peer with 5 mediocre ones… but moderate with a coverage
  // confidence multiplier so a 1-dimension specialist doesn't beat a
  // broadly-strong 4-dimension peer. Confidence = populated / total profile
  // weight; blend 50/50 with a flat floor so small coverage still scores.
  const normalized = weightedSum / weightTotal;
  const confidence = profileTotal > 0 ? weightTotal / profileTotal : 0;
  const confidenceMultiplier = 0.5 + 0.5 * confidence;
  return { score: normalized * confidenceMultiplier, coverage };
}

// ─── Provider Diversity ──────────────────────────────────────────────────────

/**
 * Interleave picks across providers so the fallback chain doesn't collapse
 * into a single provider (if that provider goes 429, every fallback fails).
 * Takes the top-N from a sorted list but skips picks whose provider already
 * appears, until we exhaust the unique providers, then cycles back.
 */
function diversifyByProvider(
  sorted: Array<{ id: string; score: number; provider: string }>,
  maxPicks: number,
): string[] {
  const picked: string[] = [];
  const seenProviders = new Set<string>();
  const stragglers: Array<{ id: string; score: number; provider: string }> = [];
  for (const m of sorted) {
    if (picked.length >= maxPicks) break;
    if (!seenProviders.has(m.provider)) {
      picked.push(m.id);
      seenProviders.add(m.provider);
    } else {
      stragglers.push(m);
    }
  }
  // Top up from stragglers in score order if we ran out of unique providers.
  for (const s of stragglers) {
    if (picked.length >= maxPicks) break;
    picked.push(s.id);
  }
  return picked;
}

// ─── Public Entry ────────────────────────────────────────────────────────────

export interface SelectOptions {
  /** Max total entries (primary + fallbacks). Default 4. */
  maxEntries?: number;
  /** Explicit benchmark data override (tests). */
  benchmarks?: BenchmarkData;
  /**
   * Ordered provider ranking for tie-breaking. When two models score
   * identically (typically same underlying weights served by different
   * endpoints), the provider earlier in this list wins. Providers not
   * listed rank after all listed providers, alphabetically.
   */
  providerPreference?: string[];
}

/**
 * Pick the best `provider/model-id` for a unit type from the candidate pool.
 * Returns null when no candidates are available.
 */
export function selectByBenchmarks(
  unitType: string,
  candidates: CandidateModel[],
  opts: SelectOptions = {},
): BenchmarkSelectionResult | null {
  if (candidates.length === 0) return null;
  const { weights, label } = profileForUnitType(unitType);
  const benchmarks = opts.benchmarks ?? loadBenchmarks();
  const maxEntries = opts.maxEntries ?? 4;

  // Build a provider-rank map. Listed providers get their index; unlisted
  // fall after all listed ones. Case-insensitive.
  const providerRank = new Map<string, number>();
  const prefList = (opts.providerPreference ?? []).map(p => p.trim().toLowerCase());
  prefList.forEach((p, i) => { if (p && !providerRank.has(p)) providerRank.set(p, i); });
  const UNLISTED_RANK = 1_000_000;
  const rankOf = (prov: string) => providerRank.get(prov) ?? UNLISTED_RANK;

  const ranked = candidates.map(c => {
    const { score, coverage } = scoreCandidate(c, weights, benchmarks);
    const fullId = `${c.provider}/${c.id}`;
    return { id: fullId, provider: c.provider.toLowerCase(), score, coverage };
  })
  // Stable sort: higher score first, then higher coverage, then
  // provider_preference rank (lower = earlier = preferred), then
  // alphabetical for determinism.
  .sort((a, b) => {
    if (b.score !== a.score) return b.score - a.score;
    if (b.coverage !== a.coverage) return b.coverage - a.coverage;
    const ra = rankOf(a.provider);
    const rb = rankOf(b.provider);
    if (ra !== rb) return ra - rb;
    return a.id.localeCompare(b.id);
  });

  const ids = diversifyByProvider(ranked, maxEntries);
  if (ids.length === 0) return null;
  const [primary, ...fallbacks] = ids;

  const scores: Record<string, number> = {};
  for (const r of ranked) scores[r.id] = Math.round(r.score * 100) / 100;
  const topCoverage = ranked[0]?.coverage ?? 0;

  return { primary, fallbacks, scores, topCoverage, profile: label };
}
