draw-things-chat

import path from "node:path";

import { MODEL_IDS, getModelFilename } from "../services/modelOverlays.js";
import {
  getCustomPresetLabelsUsingModelFilename,
  getEffectiveOverlay,
} from "../services/customConfigsLoader.js";
import {
  defaultParams as defaultTxt2Img,
  defaultParamsImg2Img as defaultImg2Img,
  defaultParamsEdit as defaultEdit,
  IMAGE_MODEL_CAPABILITY_MAP,
  MODEL_PRESET_TO_CAPABILITY_KEY,
} from "../core-bundle.mjs";
import { defaultParamsText2Video } from "../services/defaultParamsDrawThingsText2Video.js";
import { defaultParamsImage2Video } from "../services/defaultParamsDrawThingsImage2Video.js";

export type DrawThingsMode = "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid";
export type DrawThingsModeUser = "text2image" | "image2image" | "edit" | "text2video" | "image2video";

export type ResolvedImageModelPreset = {
  mode: DrawThingsModeUser;
  modelId: string;
  preset: string;
  overlaySource: "custom" | "modelOverlay" | "default";
  customConfig?: string;
};

export type ResolvedImageModelInfo = {
  modelUsedBasename: string;
  display: string;
  chosen_model_id?: string;
  model_family?: "z-image" | "qwen-image" | "flux" | "ltx" | "custom" | "unknown";
  /**
   * All known preset entries (built-in + custom overlay priority) that would
   * result in using this model file.
   */
  presets: ResolvedImageModelPreset[];
  /** Convenience: list of custom preset labels among `presets` */
  custom_config_labels?: string[];
};

function normalizeModelFilenameForLookup(value: string): string {
  return path.basename(value).trim().toLowerCase();
}

function resolveModelFamilyFromModelBasename(
  modelBasename: string
): ResolvedImageModelInfo["model_family"] {
  const normalized = normalizeModelFilenameForLookup(modelBasename);
  if (!normalized) return "unknown";

  // 1) Primary: overlay param files (z-image/qwen-image/flux/ltx)
  for (const family of ["z-image", "qwen-image", "flux", "ltx"] as const) {
    for (const mode of ["txt2img", "img2img", "edit", "txt2vid", "img2vid"] as const) {
      const m = getModelFilename(family, mode);
      if (typeof m !== "string") continue;
      if (normalizeModelFilenameForLookup(m) === normalized) return family;
    }
  }

  // 2) Secondary: capabilities.ts (older / inactive models)
  for (const family of ["z-image", "qwen-image", "flux", "ltx"] as const) {
    const key = (MODEL_PRESET_TO_CAPABILITY_KEY as any)?.[family];
    if (typeof key !== "string") continue;
    const cap = (IMAGE_MODEL_CAPABILITY_MAP as any)?.[key];
    const filenames = cap?.modelFilenames;
    if (!Array.isArray(filenames)) continue;
    for (const f of filenames) {
      if (typeof f !== "string") continue;
      if (normalizeModelFilenameForLookup(f) === normalized) return family;
    }
  }

  // 3) Tertiary: custom_configs.json presets that explicitly set `model`.
  // Only treat these as "custom" when the model file is otherwise unknown.
  const customLabels = getCustomPresetLabelsUsingModelFilename(normalized);
  if (customLabels.length > 0) return "custom";

  return "unknown";
}

function modeToUserMode(mode: DrawThingsMode): DrawThingsModeUser {
  switch (mode) {
    case "txt2img": return "text2image";
    case "img2img": return "image2image";
    case "txt2vid": return "text2video";
    case "img2vid": return "image2video";
    case "edit": return "edit";
  }
}

function userModeToMode(mode: DrawThingsModeUser): DrawThingsMode {
  switch (mode) {
    case "text2image": return "txt2img";
    case "image2image": return "img2img";
    case "text2video": return "txt2vid";
    case "image2video": return "img2vid";
    case "edit": return "edit";
  }
}

function getDefaultModelForMode(mode: DrawThingsMode): string {
  const raw =
    mode === "txt2img"
      ? (defaultTxt2Img as any)?.model
      : mode === "img2img"
      ? (defaultImg2Img as any)?.model
      : mode === "txt2vid"
      ? (defaultParamsText2Video as any)?.model
      : mode === "img2vid"
      ? (defaultParamsImage2Video as any)?.model
      : (defaultEdit as any)?.model;
  return typeof raw === "string" ? raw : "";
}

function computeEffectiveModelFilename(modelId: string, mode: DrawThingsMode): {
  modelBasename: string;
  overlaySource: "custom" | "modelOverlay" | "default";
  presetName?: string;
} {
  const base = getDefaultModelForMode(mode);
  const baseBase = normalizeModelFilenameForLookup(base);

  const overlay = getEffectiveOverlay(modelId, mode);
  const overlayModel = (overlay.params as any)?.model;
  const chosen = typeof overlayModel === "string" && overlayModel.trim()
    ? overlayModel
    : base;
  const chosenBase = normalizeModelFilenameForLookup(chosen);

  return {
    modelBasename: chosenBase || baseBase,
    overlaySource: overlay.source,
    presetName: overlay.presetName,
  };
}

function scorePresetForDisplay(p: ResolvedImageModelPreset): number {
  // Prefer the thing the user can actually pass as `model=`.
  // If a custom config overrides "auto", that should win.
  const isAuto = p.modelId === "auto";
  const isCustomConfig = p.overlaySource === "custom";
  const isOverlay = p.overlaySource === "modelOverlay";
  // Higher is better.
  return (isCustomConfig ? 100 : isOverlay ? 50 : 0) + (isAuto ? 10 : 0);
}

function pickDisplayPreset(presets: ResolvedImageModelPreset[]): ResolvedImageModelPreset | undefined {
  if (!presets.length) return undefined;
  return [...presets].sort((a, b) => scorePresetForDisplay(b) - scorePresetForDisplay(a))[0];
}

/**
 * Resolve a Draw Things backend model filename (.ckpt) to the set of presets
 * (mode + modelId) that would result in this model file being used.
 *
 * IMPORTANT: Display-only. Do not use for input validation.
 */
export function resolveImageModelInfoFromModelUsed(
  modelFilenameOrPath: string,
  opts?: { mode?: DrawThingsModeUser }
): ResolvedImageModelInfo | null {
  if (typeof modelFilenameOrPath !== "string") return null;
  const base = path.basename(modelFilenameOrPath).trim();
  if (!base) return null;

  const normalized = normalizeModelFilenameForLookup(base);
  const modes: DrawThingsMode[] = opts?.mode
    ? [userModeToMode(opts.mode)]
    : ["txt2img", "img2img", "edit", "txt2vid", "img2vid"];

  const presets: ResolvedImageModelPreset[] = [];

  for (const mode of modes) {
    for (const modelId of MODEL_IDS) {
      const eff = computeEffectiveModelFilename(modelId, mode);
      if (!eff.modelBasename) continue;
      if (eff.modelBasename !== normalized) continue;

      // Skip "default" fallback matches for any modelId other than "auto":
      // those modelIds have no overlay for this mode and silently fall through
      // to the mode default — not a meaningful/intentional preset match.
      if (eff.overlaySource === "default" && modelId !== "auto") continue;

      const userMode = modeToUserMode(mode);
      const preset = `${userMode}.${modelId}`;
      presets.push({
        mode: userMode,
        modelId,
        preset,
        overlaySource: eff.overlaySource,
        ...(eff.presetName ? { customConfig: eff.presetName } : {}),
      });
    }
  }

  const customConfigLabels = getCustomPresetLabelsUsingModelFilename(base);

  const chosen = pickDisplayPreset(presets);
  const modelFamily = resolveModelFamilyFromModelBasename(base);
  const display = modelFamily && modelFamily !== "unknown" ? `${modelFamily} (${base})` : base;

  return {
    modelUsedBasename: base,
    display,
    ...(chosen ? { chosen_model_id: chosen.modelId } : {}),
    ...(modelFamily ? { model_family: modelFamily } : {}),
    presets,
    ...(customConfigLabels.length ? { custom_config_labels: customConfigLabels } : {}),
  };
}

export function formatImageModelForToolDisplay(modelFilenameOrPath: string): string {
  const info = resolveImageModelInfoFromModelUsed(modelFilenameOrPath);
  return info?.display ?? "";
}