import type { LLamaChatPromptOptions } from "node-llama-cpp";

export const samplingProfiles = ["creative", "model-default"] as const;
export type SamplingProfile = (typeof samplingProfiles)[number];
export type GenerationKind = "post" | "reply";

type SamplingOptions = Pick<
  LLamaChatPromptOptions,
  "temperature" | "minP" | "topK" | "topP" | "repeatPenalty"
>;

export const isSamplingProfile = (value: string): value is SamplingProfile =>
  samplingProfiles.some((profile) => profile === value);

const creativeOptions = (kind: GenerationKind): SamplingOptions => ({
  temperature: kind === "post" ? 1.25 : 0.8,
  minP: kind === "post" ? 0.05 : 0.1,
  repeatPenalty: { lastTokens: 128, penalty: 1.15 },
});

/**
 * Sampling values published on each model card. These are useful as a fair
 * baseline, but are not necessarily the most interesting settings for this bot.
 */
export function getSamplingOptions(
  modelName: string,
  profile: SamplingProfile,
  kind: GenerationKind,
): SamplingOptions {
  if (profile === "creative") return creativeOptions(kind);

  const normalized = modelName.toLowerCase();
  if (normalized.includes("lfm2.5")) {
    return {
      temperature: 0.1,
      topK: 50,
      repeatPenalty: { lastTokens: 128, penalty: 1.1 },
    };
  }
  if (normalized.includes("lfm2-")) {
    return {
      temperature: 0.3,
      minP: 0.15,
      repeatPenalty: { lastTokens: 128, penalty: 1.05 },
    };
  }
  if (normalized.includes("qwen3.5")) {
    return {
      temperature: 1,
      topK: 20,
      topP: 0.95,
      // The model card's 1.5 presence penalty is outside node-llama-cpp's
      // documented 0..1 range, so use the largest supported value.
      repeatPenalty: {
        lastTokens: 128,
        penalty: 1,
        presencePenalty: 1,
      },
    };
  }
  if (normalized.includes("gemma-3") || normalized.includes("gemma-4")) {
    return {
      temperature: 1,
      topK: 64,
      topP: 0.95,
      repeatPenalty: { lastTokens: 128, penalty: 1.1 },
    };
  }

  return creativeOptions(kind);
}