draw-things-index

import path from "node:path";

export type DrawThingsMode = "text2image" | "image2image" | "edit" | "text2video" | "image2video";

export type ModelUseHints = {
  model_used: string;
  model_family: string | "unknown";
  model_use_by_mode: Partial<Record<DrawThingsMode, string[]>>;
};

export type DtcModelMappingSnapshotV1 = {
  schema?: string;
  generated_at?: string;
  models?: Record<
    string,
    {
      model_family?: string;
      preset_names?: string[];
      model_use_by_mode?: Partial<Record<DrawThingsMode, string[]>>;
    }
  >;
};

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 normalizeModelFilenameForLookup(value: string): string {
  return path.basename(value).trim().toLowerCase();
}

function tryParseInlineJsonObject(input: string): { jsonText?: string; rest: string } {
  const s = input.trimStart();
  if (!s.startsWith("{")) return { rest: input };

  let depth = 0;
  let inString = false;
  let escaped = false;

  for (let i = 0; i < s.length; i++) {
    const ch = s[i];

    if (inString) {
      if (escaped) {
        escaped = false;
        continue;
      }
      if (ch === "\\") {
        escaped = true;
        continue;
      }
      if (ch === '"') {
        inString = false;
      }
      continue;
    }

    if (ch === '"') {
      inString = true;
      continue;
    }

    if (ch === "{") {
      depth++;
      continue;
    }

    if (ch === "}") {
      depth--;
      if (depth === 0) {
        const jsonText = s.slice(0, i + 1);
        const rest = s.slice(i + 1);
        return { jsonText, rest };
      }
    }
  }

  return { rest: input };
}

function isDtcModelMappingSnapshotV1(value: unknown): value is DtcModelMappingSnapshotV1 {
  if (!value || typeof value !== "object") return false;
  const v = value as any;
  if (v.schema !== "dtc.model-mapping-snapshot.v1") return false;
  if (!v.models || typeof v.models !== "object") return false;
  return true;
}

export function parseDtcQuery(rawQuery: string): {
  query: string;
  enableModelRewrite: boolean;
  snapshot?: DtcModelMappingSnapshotV1;
} {
  const raw = String(rawQuery ?? "");
  let s = raw.trimStart();

  // New wire format: `{...snapshot...} <userQuery>` with no opt-in token.
  const { jsonText, rest } = tryParseInlineJsonObject(s);
  if (!jsonText) {
    return { query: s, enableModelRewrite: false };
  }

  try {
    const parsed = JSON.parse(jsonText);
    if (!isDtcModelMappingSnapshotV1(parsed)) {
      // Not our snapshot payload → treat as normal query, do not strip.
      return { query: s, enableModelRewrite: false };
    }
    return { query: rest.trimStart(), enableModelRewrite: true, snapshot: parsed };
  } catch {
    // Invalid JSON → treat as normal query.
    return { query: s, enableModelRewrite: false };
  }
}

export function maybeFormatModelForConsumerFromSnapshot(
  rawModel: string,
  enabled: boolean,
  snapshot?: DtcModelMappingSnapshotV1
): string {
  if (!enabled) return rawModel;

  const normalized = normalizeModelFilenameForLookup(rawModel);
  const entry = snapshot?.models?.[normalized];
  if (!entry) return rawModel;

  const family = entry.model_family && entry.model_family !== "unknown" ? entry.model_family : undefined;
  const display = family ? `${family} (${normalized})` : normalized;

  const presets: string[] = [];
  if (entry.model_use_by_mode && typeof entry.model_use_by_mode === "object") {
    for (const [mode, values] of Object.entries(entry.model_use_by_mode)) {
      if (Array.isArray(values)) {
        for (const v of values) {
          presets.push(`${mode}.${v}`);
        }
      }
    }
    presets.sort();
  }
  if (!presets.length) return display;

  return `${display}\n**Available presets for this model:** ${presets.join(", ")}`;
}

export function maybeGetModelRewriteHintsFromSnapshot(
  rawModel: string,
  enabled: boolean,
  snapshot?: DtcModelMappingSnapshotV1
): ModelUseHints | undefined {
  if (!enabled) return undefined;

  const normalized = normalizeModelFilenameForLookup(rawModel);
  const entry = snapshot?.models?.[normalized];
  if (!entry) return undefined;

  const model_family = (entry.model_family && entry.model_family.length > 0) ? entry.model_family : "unknown";

  if (entry.model_use_by_mode && typeof entry.model_use_by_mode === "object") {
    const model_use_by_mode: Partial<Record<DrawThingsMode, string[]>> = {};
    for (const mode of ["text2image", "image2image", "edit", "text2video", "image2video"] as const) {
      const values = (entry.model_use_by_mode as any)?.[mode];
      if (Array.isArray(values) && values.length) {
        model_use_by_mode[mode] = preferAutoFirst(values.map(String));
      }
    }

    return {
      model_used: normalized,
      model_family: model_family as any,
      model_use_by_mode,
    };
  }

  const model_use_by_mode: Partial<Record<DrawThingsMode, string[]>> = {};
  const presets = Array.isArray(entry.preset_names) ? entry.preset_names : [];
  for (const p of presets) {
    if (typeof p !== "string") continue;
    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]);
  }

  return {
    model_used: normalized,
    model_family: model_family as any,
    model_use_by_mode,
  };
}