draw-things-chat

import path from "node:path";

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

export type ModeInternal = "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid";
export type ModeUser = "text2image" | "image2image" | "edit" | "text2video" | "image2video";

type SnapshotPresetInfo = {
  mode: ModeUser;
  model_id: string;
  preset: string;
  model_filename: string;
  model_family: "z-image" | "qwen-image" | "flux" | "ltx" | "custom" | "unknown";
};

export type SnapshotModelInfo = {
  model_family: "z-image" | "qwen-image" | "flux" | "ltx" | "custom" | "unknown";
  model_use_by_mode: Partial<Record<ModeUser, string[]>>;
};

export type DtcModelMappingSnapshotV1 = {
  schema: "dtc.model-mapping-snapshot.v1";
  models: Record<string, SnapshotModelInfo>;
};

function toUserMode(mode: ModeInternal): ModeUser {
  switch (mode) {
    case "txt2img": return "text2image";
    case "img2img": return "image2image";
    case "edit":    return "edit";
    case "txt2vid": return "text2video";
    case "img2vid": return "image2video";
  }
}

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

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

function uniqueStrings(values: string[]): string[] {
  return [...new Set(values.filter(Boolean))];
}

function preferAutoFirst(values: string[]): string[] {
  const deduped = uniqueStrings(values);
  const rest = deduped.filter((v) => v !== "auto");
  return deduped.includes("auto") ? ["auto", ...rest] : rest;
}

function computeFamilyIndex(
  effectiveByPreset: Map<string, SnapshotPresetInfo>,
  customPresets: Map<string, any>
): Map<string, SnapshotPresetInfo["model_family"]> {
  const familyByModelBasename = new Map<string, SnapshotPresetInfo["model_family"]>();

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

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

  // 3) Tertiary: models referenced by custom presets (only if not already mapped)
  for (const preset of customPresets.values()) {
    const model = (preset as any)?.params?.model;
    if (typeof model !== "string") continue;
    const base = basenameNormalized(model);
    if (!base) continue;
    if (!familyByModelBasename.has(base)) familyByModelBasename.set(base, "custom");
  }

  // Fill for every model we actually use.
  const out = new Map<string, SnapshotPresetInfo["model_family"]>();
  for (const info of effectiveByPreset.values()) {
    out.set(info.model_filename, familyByModelBasename.get(info.model_filename) ?? "unknown");
  }

  return out;
}

function computeEffectiveModelForPreset(modelId: string, modeInternal: ModeInternal): string {
  const baseModel = getDefaultModelForMode(modeInternal);
  const baseBase = basenameNormalized(baseModel);

  const notation = `${toUserMode(modeInternal)}.${modelId}`;
  const customPreset = getCustomPreset(notation);
  const hasCustomPreset = !!customPreset;
  const customModelRaw = (customPreset as any)?.params?.model;
  const customBase =
    typeof customModelRaw === "string" && customModelRaw.trim()
      ? basenameNormalized(customModelRaw)
      : undefined;

  if (!hasCustomPreset && modelId !== "auto" && modelId !== "custom") {
    const support = checkModeSupport(modelId, toUserMode(modeInternal));
    if (!support.supported) return "";
  }

  if (modelId === "custom" && !customPreset) return "";

  const overlay = getEffectiveOverlay(modelId, modeInternal);
  const overlayModel = (overlay.params as any)?.model;
  const overlayBase =
    typeof overlayModel === "string" && overlayModel.trim()
      ? basenameNormalized(overlayModel)
      : undefined;

  return customBase ?? overlayBase ?? baseBase;
}

export async function buildDtcModelMappingSnapshot(): Promise<DtcModelMappingSnapshotV1> {
  const customPresets = await loadCustomConfigs();

  const modes: ModeInternal[] = ["txt2img", "img2img", "edit", "txt2vid", "img2vid"];

  const effectiveByPreset = new Map<string, SnapshotPresetInfo>();

  for (const modeInternal of modes) {
    const mode = toUserMode(modeInternal);

    // Use MODEL_FAMILIES (user-facing families) for snapshot, not MODEL_IDS (which includes sub-categories like qwen-edit)
    for (const modelId of MODEL_FAMILIES) {
      const preset = `${mode}.${modelId}`;
      effectiveByPreset.set(preset, {
        mode,
        model_id: modelId,
        preset,
        model_filename: computeEffectiveModelForPreset(modelId, modeInternal),
        model_family: "unknown",
      });
    }
  }

  const familyIndex = computeFamilyIndex(effectiveByPreset, customPresets as any);
  for (const info of effectiveByPreset.values()) {
    info.model_family = familyIndex.get(info.model_filename) ?? "unknown";
  }

  const presetsByModel = new Map<string, Set<string>>();
  for (const info of effectiveByPreset.values()) {
    if (!info.model_filename) continue;
    const set = presetsByModel.get(info.model_filename) ?? new Set<string>();
    set.add(info.preset);
    presetsByModel.set(info.model_filename, set);
  }

  const models: Record<string, SnapshotModelInfo> = {};
  for (const [modelFilename, presetSet] of presetsByModel.entries()) {
    const model_family = familyIndex.get(modelFilename) ?? "unknown";

    const model_use_by_mode: Partial<Record<ModeUser, string[]>> = {};
    for (const p of Array.from(presetSet).sort()) {
      const [mode, modelId] = p.split(".");
      if (!mode || !modelId) continue;
      if (mode !== "text2image" && mode !== "image2image" && mode !== "edit" && mode !== "text2video" && mode !== "image2video") continue;

      const existing = model_use_by_mode[mode] ?? [];
      model_use_by_mode[mode] = preferAutoFirst([...existing, modelId]);
    }

    models[modelFilename] = {
      model_family,
      model_use_by_mode,
    };
  }

  return {
    schema: "dtc.model-mapping-snapshot.v1",
    models,
  };
}