import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import { resolveModel } from "./model.ts";
import type { LimitKey, TreeLabelerSettings } from "./types.ts";

const CHARS_PER_TOKEN = 3.2;
const INPUT_SAFETY_TOKENS = 2048;

export interface ModelLimitProfile {
  model: string;
  contextWindow: number;
  modelMaxTokens: number;
  settings: TreeLabelerSettings;
}

export type ModelLimitAvailability =
  | { available: true; model: string; contextWindow: number; maxTokens: number }
  | { available: false; model: string; reason: string };

export function inspectModelLimitAvailability(ctx: ExtensionContext, reference: string): ModelLimitAvailability {
  try {
    const model = resolveModel(ctx, reference);
    if (!Number.isFinite(model.contextWindow) || model.contextWindow <= 0 || !Number.isFinite(model.maxTokens) || model.maxTokens <= 0) {
      return { available: false, model: `${model.provider}/${model.id}`, reason: "model catalogue has no valid context/output limits" };
    }
    return {
      available: true,
      model: `${model.provider}/${model.id}`,
      contextWindow: model.contextWindow,
      maxTokens: model.maxTokens,
    };
  } catch (error) {
    return { available: false, model: reference, reason: error instanceof Error ? error.message : String(error) };
  }
}

/**
 * Resolve capacity limits from Pi's live model catalogue. This works for
 * frontier and open-source models alike and stays current as catalogues change.
 */
export function applyModelAwareLimits(ctx: ExtensionContext, configured: TreeLabelerSettings): ModelLimitProfile {
  const model = resolveModel(ctx, configured.model);
  if (!configured.autoLimits) {
    return {
      model: `${model.provider}/${model.id}`,
      contextWindow: model.contextWindow,
      modelMaxTokens: model.maxTokens,
      settings: configured,
    };
  }

  const outputTokens = Math.max(256, model.maxTokens);
  const inputTokens = Math.max(2000, model.contextWindow - outputTokens - INPUT_SAFETY_TOKENS);
  const maxInputChars = Math.max(8000, Math.min(3_000_000, Math.floor(inputTokens * CHARS_PER_TOKEN)));

  return {
    model: `${model.provider}/${model.id}`,
    contextWindow: model.contextWindow,
    modelMaxTokens: model.maxTokens,
    settings: {
      ...configured,
      maxNodes: 10_000,
      excerptChars: 10_000,
      maxInputChars,
      // These remain bounded quality controls: labels should stay sparse and navigable.
      maxLabels: Math.max(12, Math.min(50, Math.floor(outputTokens / 256))),
      maxLabelChars: 120,
      outputTokens,
      timeoutMs: Math.max(configured.timeoutMs, 300_000),
    },
  };
}

export const LIMIT_LABELS: Record<LimitKey, string> = {
  maxNodes: "Max nodes",
  excerptChars: "Excerpt characters",
  maxInputChars: "Max input characters",
  maxLabels: "Max labels",
  maxLabelChars: "Max label characters",
  outputTokens: "Output tokens",
  timeoutMs: "Timeout seconds",
};

export function suggestedLimitValue(key: LimitKey, settings: TreeLabelerSettings): number | undefined {
  const current = settings[key];
  if (typeof current !== "number") return key === "outputTokens" ? 2800 : undefined;
  const ceilings: Record<LimitKey, number> = {
    maxNodes: 10_000,
    excerptChars: 10_000,
    maxInputChars: 3_000_000,
    maxLabels: 200,
    maxLabelChars: 200,
    outputTokens: 64_000,
    timeoutMs: 1_800_000,
  };
  return Math.min(ceilings[key], Math.max(current + 1, current * 2));
}