draw-things-chat

/**
 * Custom Configs Loader
 * Loads and manages user-defined presets from Draw Things app's custom_configs.json
 *
 * Priority System:
 * 1. Custom Configs (this file) - PRIORITY 1
 * 2. Model Overlays (modelOverlays.ts) - PRIORITY 2 (fallback)
 * 3. Defaults (defaultParamsDrawThings*.ts) - PRIORITY 3 (base)
 */

import fs from "fs";
import path from "path";
import os from "os";
import type { ImageGenerationParams } from "../core-bundle.mjs";
import { getModelOverlay } from "./modelOverlays.js";

/**
 * Status of custom configs availability
 */
export interface CustomConfigsStatus {
  available: boolean;
  filePath: string | null;
  error?: string;
}

/**
 * Parsed custom preset from Draw Things
 */
export interface CustomPreset {
  name: string; // e.g., "text2image.flux"
  mode: "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid";
  modelId: string; // e.g., "flux"
  params: Partial<ImageGenerationParams>;
}

/**
 * Raw preset structure from Draw Things custom_configs.json
 */
interface DrawThingsPreset {
  name: string; // e.g., "text2image.auto"
  configuration: Record<string, unknown>;
}

// Singleton cache for loaded presets
let cachedPresets: Map<string, CustomPreset> | null = null;
let configFilePath: string | null = null;

let didWarnCacheUninitialized = false;
let didLogUsingPreset = new Set<string>();

function isVerboseCustomConfigsLoggingEnabled(): boolean {
  // Keep per-call overlay logging OFF by default. This function intentionally
  // does not depend on draw-things-chat config wiring so that optional
  // consumers (e.g. draw-things-index) remain quiet.
  try {
    return (process as any)?.env?.DTC_CUSTOM_CONFIGS_VERBOSE === "1";
  } catch {
    return false;
  }
}

function getPresetsCacheOrNull(): Map<string, CustomPreset> | null {
  // If the caller never configured custom configs (no path), treat as disabled.
  // This is important for optional consumers (e.g. draw-things-index) that
  // import helpers without running draw-things-chat's startup sequence.
  if (cachedPresets === null && !configFilePath) {
    cachedPresets = new Map();
  }
  return cachedPresets;
}

/**
 * Check if custom_configs.json exists and is readable
 * Called ONCE during main() startup
 */
export async function checkCustomConfigs(
  customPath?: string
): Promise<CustomConfigsStatus> {
  try {
    // This function is intentionally controlled by LM Studio config (src/config.ts).
    // We do NOT auto-fallback to a default path here.
    cachedPresets = null;

    const rawPath = typeof customPath === "string" ? customPath.trim() : "";
    if (!rawPath) {
      configFilePath = null;
      cachedPresets = new Map();
      didWarnCacheUninitialized = false;
      return {
        available: false,
        filePath: null,
      };
    }

    // Expand tilde
    const expandedPath = rawPath.startsWith("~")
      ? path.join(os.homedir(), rawPath.slice(1))
      : rawPath;

    configFilePath = expandedPath;
    didWarnCacheUninitialized = false;

    // Check existence and readability
    await fs.promises.access(expandedPath, fs.constants.R_OK);

    return {
      available: true,
      filePath: expandedPath,
    };
  } catch (error) {
    const errorMessage =
      error instanceof Error ? error.message : "Unknown error";
    return {
      available: false,
      filePath: null,
      error: errorMessage,
    };
  }
}

/**
 * Load all custom presets from custom_configs.json
 * Returns Map<presetName, CustomPreset>
 * Caches result for performance
 */
export async function loadCustomConfigs(): Promise<Map<string, CustomPreset>> {
  // Return cached if available
  if (cachedPresets !== null) {
    return cachedPresets;
  }

  const presets = new Map<string, CustomPreset>();

  try {
    if (!configFilePath) {
      // Not configured / disabled.
      cachedPresets = presets;
      return presets;
    }

    // Read and parse JSON
    const fileContent = await fs.promises.readFile(configFilePath, "utf-8");
    const rawPresets: DrawThingsPreset[] = JSON.parse(fileContent);

    if (!Array.isArray(rawPresets)) {
      console.warn(
        "[CustomConfigs] Expected array of presets, got:",
        typeof rawPresets
      );
      cachedPresets = presets;
      return presets;
    }

    // Parse each preset
    for (const rawPreset of rawPresets) {
      try {
        const preset = parsePreset(rawPreset);
        if (preset) {
          presets.set(preset.name, preset);
        }
      } catch (error) {
        // Silently skip invalid presets (Draw Things allows freeform names)
      }
    }

    if (presets.size > 0) {
      const presetNames = Array.from(presets.keys()).join(", ");
      console.log(
        `[CustomConfigs] Loaded ${presets.size} presets: ${presetNames}`
      );
    } else {
      console.log(
        `[CustomConfigs] No valid presets found in ${configFilePath}`
      );
    }
  } catch (error) {
    console.warn(
      "[CustomConfigs] Failed to load custom_configs.json:",
      error instanceof Error ? error.message : error
    );
  }

  // Cache result (even if empty)
  cachedPresets = presets;
  didWarnCacheUninitialized = false;
  return presets;
}

/**
 * Convert Draw Things camelCase parameter names to our snake_case
 * Handles nested objects (e.g., loras array)
 */
function convertParamNames(
  config: Record<string, unknown>
): Record<string, unknown> {
  const converted: Record<string, unknown> = {};

  for (const [key, value] of Object.entries(config)) {
    // Convert camelCase to snake_case
    // Handles cases like: guidanceScale → guidance_scale, clipLText → clip_l_text
    const snakeKey = key
      .replace(/([a-z0-9])([A-Z])/g, "$1_$2")
      .replace(/([A-Z])([A-Z][a-z])/g, "$1_$2")
      .toLowerCase();

    // Handle nested arrays (e.g., loras)
    if (Array.isArray(value)) {
      converted[snakeKey] = value.map((item) =>
        typeof item === "object" && item !== null
          ? convertParamNames(item as Record<string, unknown>)
          : item
      );
    } else if (typeof value === "object" && value !== null) {
      // Handle nested objects
      converted[snakeKey] = convertParamNames(value as Record<string, unknown>);
    } else {
      converted[snakeKey] = value;
    }
  }

  return converted;
}

/**
 * Parse a single preset from Draw Things format
 * Converts camelCase configuration to snake_case
 */
function parsePreset(rawPreset: DrawThingsPreset): CustomPreset | null {
  const { name, configuration } = rawPreset;

  if (!name || typeof name !== "string") {
    console.warn("[CustomConfigs] Preset missing name field");
    return null;
  }

  // Parse name format: "mode.modelId" (e.g., "text2image.flux")
  const parts = name.split(".");
  if (parts.length !== 2) {
    // Silently skip - Draw Things allows freeform preset names
    return null;
  }

  const [modeRaw, modelId] = parts;

  // Map Draw Things mode names to our internal mode types
  const modeMap: Record<string, "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid"> = {
    text2image: "txt2img",
    image2image: "img2img",
    edit: "edit",
    text2video: "txt2vid",
    image2video: "img2vid",
  };

  const mode = modeMap[modeRaw];
  if (!mode) {
    // Silently skip - only process text2image/image2image/edit presets
    return null;
  }

  // Convert camelCase configuration to snake_case
  const convertedConfig = convertParamNames(configuration);

  // Cast to our params type (will be validated by filterCustomPresetParams in Task 1.4)
  const params = convertedConfig as Partial<ImageGenerationParams>;

  return {
    name,
    mode,
    modelId,
    params,
  };
}

/**
 * Get specific preset by notation (e.g., "text2image.flux")
 * Returns null if not found
 */
export function getCustomPreset(notation: string): CustomPreset | null {
  const cache = getPresetsCacheOrNull();
  if (!cache) {
    if (!didWarnCacheUninitialized) {
      console.warn(
        "[CustomConfigs] Cache not initialized, call loadCustomConfigs() first"
      );
      didWarnCacheUninitialized = true;
    }
    return null;
  }

  return cache.get(notation) ?? null;
}

/**
 * Clear cache (for testing or hot-reload scenarios)
 */
export function clearCustomConfigsCache(): void {
  cachedPresets = null;
  configFilePath = null;
  didWarnCacheUninitialized = false;
  didLogUsingPreset = new Set();
}

/**
 * Get effective overlay for mode + model combination
 *
 * Priority:
 * 1. Custom Config (if available) - ALWAYS checked first
 * 2. Model Overlay (existing system)
 * 3. null (use defaults)
 *
 * IMPORTANT: Custom Configs are checked even when modelId is undefined.
 * In that case, we default to "auto" as the effective model identifier.
 * This ensures that Custom Configs like "text2image.auto" are always used
 * when available, even if the user doesn't explicitly pass model="auto".
 */
export function getEffectiveOverlay(
  modelId: string | undefined,
  mode: "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid"
): {
  source: "custom" | "modelOverlay" | "default";
  presetName?: string;
  params: Partial<ImageGenerationParams> | null;
} {
  // Step 1: Check custom configs (PRIORITY 1)
  // ALWAYS check Custom Configs, even if modelId is undefined!
  // Default to "auto" as the effective model identifier.
  const cache = getPresetsCacheOrNull();
  if (cache) {
    const effectiveModelId = modelId ?? "auto";
    const notation = buildPresetNotation(mode, effectiveModelId);
    const verbose = isVerboseCustomConfigsLoggingEnabled();
    if (verbose) {
      console.debug(
        `[CustomConfigs] getEffectiveOverlay: modelId=${effectiveModelId}, mode=${mode}, notation=${notation}, cacheSize=${cache.size}`
      );
    }
    const customPreset = getCustomPreset(notation);
    if (customPreset) {
      // Avoid spamming logs when called repeatedly (e.g. metadata resolver loops).
      if (verbose && !didLogUsingPreset.has(notation)) {
        console.info(`[CustomConfigs] Using custom preset '${notation}'`);
        didLogUsingPreset.add(notation);
      }
      return {
        source: "custom",
        presetName: notation,
        params: customPreset.params,
      };
    } else {
      if (verbose) {
        console.debug(
          `[CustomConfigs] Custom preset '${notation}' not found, falling back to model overlay`
        );
      }
    }
  } else {
    // Only warn if a config path exists (i.e. the feature is enabled), and de-dupe spam.
    if (configFilePath && !didWarnCacheUninitialized) {
      console.warn(
        `[CustomConfigs] cachedPresets is null (loadCustomConfigs() was not called or failed)`
      );
      didWarnCacheUninitialized = true;
    }
  }

  // Step 2: Fall back to model overlays (PRIORITY 2)
  const modelOverlay = getModelOverlay(modelId, mode);
  if (modelOverlay) {
    return {
      source: "modelOverlay",
      params: modelOverlay,
    };
  }

  // Step 3: Use defaults
  return {
    source: "default",
    params: null,
  };
}

/**
 * Build preset notation from mode and modelId
 * Examples:
 *   txt2img + flux → "text2image.flux"
 *   img2img + custom → "image2image.custom"
 *   edit + z-image → "edit.z-image"
 */
function buildPresetNotation(
  mode: "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid",
  modelId: string
): string {
  const modeMap: Record<string, string> = {
    txt2img: "text2image",
    img2img: "image2image",
    edit: "edit",
    txt2vid: "text2video",
    img2vid: "image2video",
  };
  return `${modeMap[mode]}.${modelId}`;
}

/**
 * Get all available mode+model combinations from Custom Configs
 * Returns array of tuples [mode, modelId]
 * Used for error handling to show what's available
 */
export function getAvailableCustomCombinations(): Array<
  ["text2image" | "image2image" | "edit" | "text2video" | "image2video", string]
> {
  const cache = getPresetsCacheOrNull();
  if (!cache) return [];

  const combinations: Array<["text2image" | "image2image" | "edit" | "text2video" | "image2video", string]> =
    [];
  for (const preset of cache.values()) {
    // Map internal mode back to user-facing mode
    const userMode =
      preset.mode === "txt2img"
        ? "text2image"
        : preset.mode === "img2img"
        ? "image2image"
        : preset.mode === "txt2vid"
        ? "text2video"
        : preset.mode === "img2vid"
        ? "image2video"
        : "edit";
    combinations.push([userMode, preset.modelId]);
  }

  return combinations;
}

/**
 * Return the Draw Things custom preset names that explicitly reference a given model filename.
 *
 * NOTE: Many presets do not set `model` explicitly; those cannot be matched here.
 * This is intended for optional, display-only enrichment.
 */
export function getCustomPresetLabelsUsingModelFilename(
  modelFilenameOrPath: string
): string[] {
  const cache = getPresetsCacheOrNull();
  if (!cache) return [];
  if (typeof modelFilenameOrPath !== "string") return [];

  const wanted = path.basename(modelFilenameOrPath).trim().toLowerCase();
  if (!wanted) return [];

  const labels = new Set<string>();
  for (const preset of cache.values()) {
    const presetModel = (preset.params as any)?.model;
    if (typeof presetModel !== "string") continue;
    const presetBase = path.basename(presetModel).trim().toLowerCase();
    if (presetBase === wanted) {
      labels.add(preset.name);
    }
  }

  return Array.from(labels).sort();
}