import type { ThinkingLevel, ThinkingLevelMap } from "@earendil-works/pi-ai";
import { isGpt55Model, isGpt6Model } from "./model-api.ts";
import type { ModelsDevMetadata } from "./types.ts";

export interface ModelCapabilityContext {
  availableModelId: string;
  metadataModelId?: string;
}

export interface ModelCapabilityOverrides {
  reasoning?: boolean;
  thinkingLevelMap?: ThinkingLevelMap;
}

interface ModelCapabilityRule {
  matches: (context: ModelCapabilityContext) => boolean;
  overrides: ModelCapabilityOverrides;
}

/**
 * Fallback for GPT-5.5 when the running pi's `openai-codex` catalog no longer
 * lists it. Mirrors pi's native definition: no `minimal` on the Codex route,
 * `xhigh` but no `max`.
 */
const GPT_5_5_THINKING_LEVEL_MAP: ThinkingLevelMap = {
  minimal: "low",
  xhigh: "xhigh",
};

/**
 * Fallback for GPT-5.6 models newer than the running pi's `openai-codex`
 * catalog, whose native map wins when it knows the model. The Codex route has
 * no `minimal` effort: CLIProxyAPI rejects it (`level "minimal" not supported`),
 * so pi's `minimal` maps up to `low`, as pi's native definitions do.
 */
const GPT_5_6_THINKING_LEVEL_MAP: ThinkingLevelMap = {
  off: "none",
  minimal: "low",
  low: "low",
  medium: "medium",
  high: "high",
  xhigh: "xhigh",
  max: "max",
};

/**
 * Fallback for GPT-6 models newer than the running pi's `openai-codex`
 * catalog. GPT-6 Astra cannot switch reasoning off (its catalog entry lists only
 * low..max plus an `ultra` level pi has no slot for), and it has no `minimal`
 * effort: the lowest it accepts is `low`, so pi's `minimal` maps down to it.
 * Mirrors pi's native `openai-codex/gpt-6-astra` definition.
 */
const GPT_6_THINKING_LEVEL_MAP: ThinkingLevelMap = {
  off: null,
  minimal: "low",
  low: "low",
  medium: "medium",
  high: "high",
  xhigh: "xhigh",
  max: "max",
};

function includesModelFamily(context: ModelCapabilityContext, family: string): boolean {
  return [context.availableModelId, context.metadataModelId]
    .filter((id): id is string => id !== undefined)
    .some((id) => id.includes(family));
}

const MODEL_CAPABILITY_RULES: readonly ModelCapabilityRule[] = [
  {
    matches: (context) => isGpt55Model(context),
    overrides: {
      reasoning: true,
      thinkingLevelMap: GPT_5_5_THINKING_LEVEL_MAP,
    },
  },
  {
    matches: (context) => includesModelFamily(context, "gpt-5.6"),
    overrides: {
      reasoning: true,
      thinkingLevelMap: GPT_5_6_THINKING_LEVEL_MAP,
    },
  },
  {
    matches: (context) => isGpt6Model(context),
    overrides: {
      reasoning: true,
      thinkingLevelMap: GPT_6_THINKING_LEVEL_MAP,
    },
  },
];

/**
 * Pi's non-`off` thinking levels in rank order. Each is also the literal effort
 * name Anthropic and OpenAI accept, so a models.dev `effort` list can be read
 * directly against it.
 */
const PI_EFFORT_LEVELS: readonly ThinkingLevel[] = [
  "minimal",
  "low",
  "medium",
  "high",
  "xhigh",
  "max",
];

/**
 * A model's thinking map: published verbatim when the metadata carries one,
 * otherwise derived from models.dev `reasoning_options`.
 *
 * Pi's built-in catalog publishes a finished map per model and is the seed this
 * package starts from, so that map is used as-is. models.dev never carries one,
 * which is what the derivation below exists for.
 *
 * Pi treats an omitted map as "standard levels through `high`, no `xhigh` or
 * `max`", so a model that accepts those efforts (Claude Fable 5.x, Opus 4.7+,
 * Sonnet 5, GPT-5.6) silently loses its top levels unless something publishes
 * them. CLIProxyAPI's `/v1/models` does not, but the models.dev catalog the
 * extension already fetches lists the exact effort names each model accepts.
 *
 * Only the `effort` option contributes. Each pi level named in it maps to
 * itself; every other level maps to `null`, which is how pi represents holes
 * (Opus 4.6 accepts `max` but not `xhigh`). A `none` effort maps to `off`.
 * When `effort` is the *only* option, thinking cannot be disabled at all (Claude
 * Fable 5.x rejects `thinking.type = disabled`), so `off` becomes `null`; a
 * `toggle` or `budget_tokens` sibling means the model also has a switchable
 * mode, and `off` stays available. Models that only expose `budget_tokens`
 * return `undefined` so pi keeps its default budget mapping.
 */
export function thinkingLevelMapFromMetadata(metadata: ModelsDevMetadata): ThinkingLevelMap | undefined {
  if (metadata.thinkingLevelMap) return metadata.thinkingLevelMap;

  const options = metadata.reasoning_options;
  if (!Array.isArray(options)) return undefined;

  const effort = options.find((option) => option?.type === "effort" && Array.isArray(option.values));
  if (!effort?.values) return undefined;

  const map = mapFromAcceptedEfforts(new Set(effort.values));
  if (!map.off && options.every((option) => option?.type === "effort")) map.off = null;
  return map;
}

function mapFromAcceptedEfforts(accepted: ReadonlySet<string>): ThinkingLevelMap {
  const map: ThinkingLevelMap = {};
  for (const level of PI_EFFORT_LEVELS) {
    map[level] = accepted.has(level) ? level : null;
  }
  if (accepted.has("none")) map.off = "none";
  return map;
}

/**
 * A thinking map from the effort names CLIProxyAPI lists for a model, for
 * models no better source describes (Devin's own SWE models, for example).
 *
 * Each listed pi level maps to itself and every other level to `null`; `none`
 * maps to `off`. Unlike models.dev, CPA's list cannot say that thinking is
 * impossible to switch off, so without `none` the `off` slot is left to pi's
 * default. Returns `undefined` when no listed name is a pi level.
 */
export function thinkingLevelMapFromEfforts(efforts: readonly string[]): ThinkingLevelMap | undefined {
  const accepted = new Set(efforts.map((effort) => effort.toLowerCase()));
  if (!PI_EFFORT_LEVELS.some((level) => accepted.has(level))) return undefined;
  return mapFromAcceptedEfforts(accepted);
}

export function getModelCapabilityOverrides(context: ModelCapabilityContext): ModelCapabilityOverrides {
  const resolved: ModelCapabilityOverrides = {};

  for (const rule of MODEL_CAPABILITY_RULES) {
    if (!rule.matches(context)) continue;
    if (rule.overrides.reasoning !== undefined) resolved.reasoning = rule.overrides.reasoning;
    if (rule.overrides.thinkingLevelMap) resolved.thinkingLevelMap = { ...rule.overrides.thinkingLevelMap };
  }

  return resolved;
}