draw-things-chat

import {
  defaultParams,
  defaultParamsImg2Img,
  drawthingsLimits,
  defaultParamsEdit,
  drawthingsEditLimits,
  ImageGenerationParams,
  validateImageGenerationParams,
  getSize as imgGetSize,
  toPng as imgToPng,
  DrawThingsGenerationResult,
} from "../core-bundle.mjs";
import { ModelId } from "./modelOverlays.js";
import { getEffectiveOverlay } from "./customConfigsLoader.js";
import axios, { AxiosInstance } from "axios";
import fs from "fs";
import path from "path";
import { fileURLToPath } from "url";
import { ImageBackend, BackendGenerateResult, ProgressCallback } from "./imageBackend.js";

// HTTP expects exact EnumParameter.commandLineAbbreviation strings for some enums.
// Custom configs come from Draw Things and may contain numeric FlatBuffer enum codes.
const HTTP_SEED_MODE_FROM_CODE: Record<number, string> = {
  0: "Legacy",
  1: "Torch CPU Compatible",
  2: "Scale Alike",
  3: "NVIDIA GPU Compatible",
};

const HTTP_SEED_MODE_FROM_KEY: Record<string, string> = {
  Legacy: "Legacy",
  TorchCpuCompatible: "Torch CPU Compatible",
  TorchCPUCompatible: "Torch CPU Compatible",
  ScaleAlike: "Scale Alike",
  NvidiaGpuCompatible: "NVIDIA GPU Compatible",
  NvidiaGPUCompatible: "NVIDIA GPU Compatible",
};

const HTTP_SAMPLER_FROM_CODE: Record<number, string> = {
  0: "DPM++ 2M Karras",
  1: "Euler a",
  2: "DDIM",
  3: "PLMS",
  4: "DPM++ SDE Karras",
  5: "UniPC",
  6: "LCM",
  7: "Euler A Substep",
  8: "DPM++ SDE Substep",
  9: "TCD",
  10: "Euler A Trailing",
  11: "DPM++ SDE Trailing",
  12: "DPM++ 2M AYS",
  13: "Euler A AYS",
  14: "DPM++ SDE AYS",
  15: "DPM++ 2M Trailing",
  16: "DDIM Trailing",
  17: "UniPC Trailing",
  18: "UniPC AYS",
  19: "TCD Trailing",
};

const HTTP_SAMPLER_FROM_KEY: Record<string, string> = {
  DPMPP2MKarras: "DPM++ 2M Karras",
  EulerA: "Euler a",
  DDIM: "DDIM",
  PLMS: "PLMS",
  DPMPPSDEKarras: "DPM++ SDE Karras",
  UniPC: "UniPC",
  LCM: "LCM",
  EulerASubstep: "Euler A Substep",
  DPMPPSDESubstep: "DPM++ SDE Substep",
  TCD: "TCD",
  EulerATrailing: "Euler A Trailing",
  DPMPPSDETrailing: "DPM++ SDE Trailing",
  DPMPP2MAYS: "DPM++ 2M AYS",
  EulerAAYS: "Euler A AYS",
  DPMPPSDEAYS: "DPM++ SDE AYS",
  DPMPP2MTrailing: "DPM++ 2M Trailing",
  DDIMTrailing: "DDIM Trailing",
  UniPCTrailing: "UniPC Trailing",
  UniPCAYS: "UniPC AYS",
  TCDTrailing: "TCD Trailing",
};

function normalizeHttpSeedMode(value: unknown): unknown {
  if (typeof value === "number" && Number.isFinite(value)) {
    return HTTP_SEED_MODE_FROM_CODE[value] ?? String(value);
  }
  if (typeof value === "string" && value.trim()) {
    const raw = value.trim();
    // If already a spaced HTTP string, keep it.
    if (raw.includes(" ")) return raw;
    // Else treat as enum key (or space-stripped key)
    const key = raw.replace(/\s+/g, "");
    return HTTP_SEED_MODE_FROM_KEY[key] ?? raw;
  }
  return value;
}

function normalizeHttpSampler(value: unknown): unknown {
  if (typeof value === "number" && Number.isFinite(value)) {
    return HTTP_SAMPLER_FROM_CODE[value] ?? String(value);
  }
  if (typeof value === "string" && value.trim()) {
    const raw = value.trim();
    // If already a spaced HTTP string, keep it.
    if (raw.includes(" ") || raw.includes("+")) return raw;
    const key = raw.replace(/\s+/g, "");
    return HTTP_SAMPLER_FROM_KEY[key] ?? raw;
  }
  return value;
}

function applyHttpCustomOverlayFixups(
  overlay: Partial<ImageGenerationParams>
): Partial<ImageGenerationParams> {
  const out: any = { ...(overlay as any) };

  // The HTTP API accepts a strict subset of config keys.
  // We only normalize/drop what must be normalized/dropped based on what the
  // HTTP defaults in defaultParamsDrawThingsImg2Img.ts use.
  const rename = (from: string, to: string) => {
    if (!Object.prototype.hasOwnProperty.call(out, from)) return;
    if (!Object.prototype.hasOwnProperty.call(out, to)) out[to] = out[from];
    delete out[from];
  };

  // stage2Guidance/Shift exist in custom configs; HTTP defaults use stage_2_*.
  rename("stage2_guidance", "stage_2_guidance");
  rename("stage2_shift", "stage_2_shift");
  // stage2_steps has no known HTTP counterpart in our defaults.
  delete out.stage2_steps;
  rename("stage2Guidance", "stage_2_guidance");
  rename("stage2Shift", "stage_2_shift");
  delete out.stage2Steps;

  // Text-encoder toggle naming differs: defaults use t5_text_encoder_decoding.
  rename("t5_text_encoder", "t5_text_encoder_decoding");
  rename("t5TextEncoder", "t5_text_encoder_decoding");

  // imageGuidanceScale is rejected by HTTP; defaults use image_guidance.
  rename("image_guidance_scale", "image_guidance");
  rename("imageGuidanceScale", "image_guidance");

  // Keys explicitly rejected by the HTTP schema.
  delete out.target_image_height;
  delete out.target_image_width;
  delete out.targetImageHeight;
  delete out.targetImageWidth;
  delete out.original_image_height;
  delete out.original_image_width;
  delete out.originalImageHeight;
  delete out.originalImageWidth;
  delete out.negative_original_image_height;
  delete out.negative_original_image_width;
  delete out.negativeOriginalImageHeight;
  delete out.negativeOriginalImageWidth;
  delete out.face_restoration;
  delete out.faceRestoration;

  // Size keys must NOT come from custom configs in HTTP mode.
  // The tool/core controls requested size, and img2img must match init_images.
  delete out.width;
  delete out.height;
  delete out.target_width;
  delete out.target_height;

  // Batch keys must NOT come from custom configs.
  // The tool/core controls batch sizing via variants parameter.
  delete out.batch_count;
  delete out.batch_size;
  delete out.batchCount;
  delete out.batchSize;

  // Upscaler keys must NOT come from custom configs.
  // The tool/core controls upscaling via _dt_needs_upscaler decision.
  delete out.upscaler;
  delete out.upscaler_scale;
  delete out.upscalerScale;

  // Filter tea cache window params when tea_cache is false.
  const teaCache =
    typeof out.tea_cache === "boolean"
      ? out.tea_cache
      : typeof out.teaCache === "boolean"
      ? out.teaCache
      : undefined;
  if (teaCache === false) {
    delete out.tea_cache_start;
    delete out.tea_cache_end;
    delete out.teaCacheStart;
    delete out.teaCacheEnd;
  }

  // HTTP backend rejects sending both of these keys at once.
  // Our defaults use `upscaler_scale`, while custom_configs may provide
  // `upscalerScaleFactor` / `upscaler_scale_factor`.
  if (Object.prototype.hasOwnProperty.call(out, "upscaler_scale_factor")) {
    if (!Object.prototype.hasOwnProperty.call(out, "upscaler_scale")) {
      out.upscaler_scale = out.upscaler_scale_factor;
    }
    delete out.upscaler_scale_factor;
  }
  if (Object.prototype.hasOwnProperty.call(out, "upscalerScaleFactor")) {
    if (!Object.prototype.hasOwnProperty.call(out, "upscaler_scale")) {
      out.upscaler_scale = out.upscalerScaleFactor;
    }
    delete out.upscalerScaleFactor;
  }
  // If any upstream mapping produced camelCase `upscalerScale`, drop it.
  if (Object.prototype.hasOwnProperty.call(out, "upscalerScale")) {
    delete out.upscalerScale;
  }

  // Convert enum-like values to HTTP strings.
  if (Object.prototype.hasOwnProperty.call(out, "seed_mode")) {
    out.seed_mode = normalizeHttpSeedMode(out.seed_mode) as any;
  }
  if (Object.prototype.hasOwnProperty.call(out, "seedMode")) {
    out.seedMode = normalizeHttpSeedMode(out.seedMode) as any;
  }
  if (Object.prototype.hasOwnProperty.call(out, "sampler")) {
    out.sampler = normalizeHttpSampler(out.sampler) as any;
  }

  return out as Partial<ImageGenerationParams>;
}

/**
 * simplified DrawThingsService
 * focus on core functionality: connect to Draw Things API and generate image
 */
type AuditMode = "txt2img" | "img2img" | "edit" | "txt2vid" | "img2vid";
function buildAuditLogger({
  backend,
  mode,
}: {
  backend: string;
  mode: AuditMode;
}) {
  function localTimestamp(): string {
    try {
      return new Date().toLocaleString(undefined, {
        year: "numeric",
        month: "2-digit",
        day: "2-digit",
        hour: "2-digit",
        minute: "2-digit",
        second: "2-digit",
        hour12: false,
        timeZoneName: "short",
      });
    } catch {
      return new Date().toString();
    }
  }
  const requestId = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
  const entry: any = {
    timestamp: localTimestamp(),
    backend,
    mode,
    requestId,
  };

  function setRequested(obj: Record<string, unknown>) {
    entry.requested = obj;
  }
  function setUserInput(obj: Record<string, unknown>) {
    entry.user_input = obj;
  }
  function setPrompt(p: string | undefined) {
    if (typeof p === "string" && p.trim().length > 0) entry.prompt = p;
  }
  function setBackendReturned(obj: Record<string, unknown>) {
    entry.backendReturned = obj;
  }
  function setPostProcessed(obj: Record<string, unknown>) {
    entry.postProcessed = obj;
  }
  function setExtras(obj: Record<string, unknown>) {
    entry.extras = { ...(entry.extras || {}), ...obj };
  }
  function setInferenceTime(ms: number) {
    entry.inference_time_ms = ms;
  }
  function setError(err: unknown) {
    let message = "unknown error";
    let status: number | undefined = undefined;
    if (typeof err === "string") message = err;
    else if (err && typeof err === "object") {
      const anyErr = err as any;
      message = anyErr.message || JSON.stringify(anyErr);
      if (typeof anyErr.status === "number") status = anyErr.status;
    }
    entry.error = status ? { message, status } : { message };
  }

  async function write() {
    try {
      // Use centralized logs directory from helpers/paths to avoid divergence
      const { getLogsDir } = await import("../core-bundle.mjs");
      const logsDir = getLogsDir();
      const filePath = path.resolve(
        logsDir,
        "generate-image-plugin.audit.jsonl"
      );
      await fs.promises.mkdir(logsDir, { recursive: true });
      const block = JSON.stringify(entry, null, 2) + "\n\n";
      await fs.promises.appendFile(filePath, block, { encoding: "utf8" });
    } catch (e) {
      console.error(
        "auditLog write failed:",
        e instanceof Error ? e.message : String(e)
      );
    }
  }

  return {
    requestId,
    setPrompt,
    setUserInput,
    setRequested,
    setBackendReturned,
    setPostProcessed,
    setExtras,
    setInferenceTime,
    setError,
    write,
  };
}

export class DrawThingsService implements ImageBackend {
  public readonly name = "drawthings";
  // make baseUrl public for compatibility with index.ts
  public baseUrl: string;
  // change to public axios for compatibility
  public axios: AxiosInstance;

  constructor(
    apiUrl = "http://127.0.0.1:7860",
    sharedSecretParam?: string | null
  ) {
    this.baseUrl = apiUrl;

    // initialize axios
    const sharedSecret =
      (sharedSecretParam ?? process.env.DRAWTHINGS_SHARED_SECRET) || undefined;
    this.axios = axios.create({
      baseURL: this.baseUrl,
      timeout: 300000, // 5 minutes timeout (image generation may take time)
      headers: {
        "Content-Type": "application/json",
        Accept: "application/json",
        "Accept-Encoding": "gzip, deflate, br",
        ...(sharedSecret ? { Authorization: `Bearer ${sharedSecret}` } : {}),
      },
    });

    // log initialization
    console.info(
      `DrawThingsService initialized, API location: ${this.baseUrl}`
    );
  }

  /**
   * Set new base URL and update axios instance
   * @param url new base URL
   */
  setBaseUrl(url: string): void {
    this.baseUrl = url;
    this.axios.defaults.baseURL = url;
    console.info(`Updated API base URL to: ${url}`);
  }

  /**
   * check API connection
   * simplified version that just checks if API is available
   */
  async checkApiConnection(): Promise<boolean> {
    try {
      console.debug(`Checking API connection to: ${this.baseUrl}`);

      // Try simple endpoint with short timeout
      const response = await this.axios.get("/sdapi/v1/options", {
        timeout: 5000,
        validateStatus: (status) => status >= 200,
      });

      const isConnected = response.status >= 200;
      console.info(
        `API connection check: ${isConnected ? "Success" : "Failed"}`
      );
      return isConnected;
    } catch (error) {
      console.error(`API connection check failed: ${(error as Error).message}`);
      return false;
    }
  }

  /**
   * get default params
   */
  getDefaultParams(): ImageGenerationParams {
    return defaultParams;
  }

  /**
   * generate image (text-to-image)
   * @param inputParams user provided params (raw from tool). Backend filters to engine-allowed keys.
   * @param _onProgress Progress callback (not supported by HTTP backend, ignored)
   */
  async generateImage(
    inputParams: Partial<ImageGenerationParams> = {},
    _onProgress?: ProgressCallback
  ): Promise<BackendGenerateResult> {
    const audit = buildAuditLogger({ backend: this.name, mode: "txt2img" });
    // Overlay tracking for result metadata
    let overlaySource: "custom" | "modelOverlay" | "default" = "default";
    let overlayPreset: string | undefined;
    let effectiveModel: string | undefined;
    let effectiveLoras: string[] | undefined;

    try {
      // Measure only backend latency (request -> response), like gRPC path
      let startTime = 0;
      let endTime = 0;
      // 1) Validate raw input shape (basic checks only)
      let params: Partial<ImageGenerationParams> = {};
      try {
        const validationResult = validateImageGenerationParams(inputParams);
        if (validationResult.valid) {
          params = inputParams;
        } else {
          console.warn("parameter validation failed, use default params");
        }
      } catch (error) {
        console.warn("parameter validation error:", error);
      }

      // 2) Special case: random_string → prompt
      if (
        params.random_string &&
        (!params.prompt || Object.keys(params).length === 1)
      ) {
        params.prompt = params.random_string;
        delete params.random_string;
      }

      // 3) Ensure prompt
      if (!params.prompt) {
        params.prompt = inputParams.prompt || defaultParams.prompt;
      }

      // 4) Canonical mapping (imageFormat, quality)
      let usedImageFormat: string | undefined = undefined;
      let usedQuality: string | undefined = undefined;
      try {
        const fmt = (inputParams as any).imageFormat as string | undefined;
        // Shorthand should not override explicit dimensions (core passes requested_effective via width/height).
        if ((params as any).width == null && (params as any).height == null) {
          if (fmt === "square") {
            params.width = 1024;
            params.height = 1024;
            usedImageFormat = "square";
          } else if (fmt === "landscape") {
            params.width = 1024;
            params.height = 768;
            usedImageFormat = "landscape";
          } else if (fmt === "portrait") {
            params.width = 768;
            params.height = 1024;
            usedImageFormat = "portrait";
          } else if (fmt === "16:9") {
            params.width = 1024;
            params.height = 576;
            usedImageFormat = "16:9";
          }
        }

        const qual = (inputParams as any).quality as string | undefined;
        if (qual === "low") {
          params.steps = 4;
          usedQuality = "low";
        } else if (qual === "medium") {
          params.steps = 8;
          usedQuality = "medium";
        } else if (qual === "high") {
          params.steps = 12;
          usedQuality = "high";
        }
        // 'auto' -> do not override steps (keep defaults)

        const variants = (inputParams as any).variants as number | undefined;
        if (typeof variants === "number") {
          const v = Math.max(1, Math.min(4, Math.round(variants)));
          (params as any).batch_size = v;
          (params as any).batch_count = 1;
        }
      } catch {}

      // 5) Filter to engine-allowed keys (based on defaultParams keys)
      const allowedKeys = new Set(Object.keys(defaultParams));
      const filteredParams: Partial<ImageGenerationParams> = {};
      for (const [k, v] of Object.entries(params)) {
        if (allowedKeys.has(k)) {
          (filteredParams as any)[k] = v as any;
        }
      }

      // 5b) Apply effective overlay (Custom Configs → Model Overlay → Defaults)
      const modelId = (inputParams as any).model as ModelId | undefined;
      const {
        source,
        presetName,
        params: overlayParams,
      } = getEffectiveOverlay(modelId, "txt2img");

      const overlayParamsHttp =
        source === "custom" && overlayParams
          ? applyHttpCustomOverlayFixups(overlayParams)
          : overlayParams;

      // Capture overlay info for result metadata
      overlaySource = source;
      overlayPreset = presetName;

      // Tool-level preset 'auto' means: do not override engine model.
      // LM Studio UI may send "auto" explicitly; never forward that literal to Draw Things.
      if (modelId === "auto") {
        delete (filteredParams as any).model;
      }
      if (overlayParamsHttp) {
        console.info(
          `Applying ${source} overlay for '${modelId ?? 'auto'}' (txt2img)${
            presetName ? ` preset=${presetName}` : ""
          }`
        );
        // Remove 'model' from filteredParams so overlay model filename wins
        delete (filteredParams as any).model;
      }

      // Log overlay source to audit
      audit.setExtras({
        overlay_source: source,
        ...(presetName && { overlay_preset: presetName }),
      });

      // Seed policy:
      // - Tool interface does not accept seed (core uses strict minimal schema).
      // - Custom Configs MAY set seed (including -1) and should be effective.
      const effectiveSeed =
        source === "custom" &&
        typeof (overlayParamsHttp as any)?.seed === "number" &&
        Number.isFinite((overlayParamsHttp as any)?.seed)
          ? (overlayParamsHttp as any).seed
          : defaultParams.seed;

      // Seed mode policy mirrors seed: allow Custom Configs to override; otherwise keep defaults.
      const effectiveSeedMode =
        source === "custom" &&
        typeof (overlayParamsHttp as any)?.seed_mode === "string" &&
        String((overlayParamsHttp as any).seed_mode).trim()
          ? String((overlayParamsHttp as any).seed_mode)
          : defaultParams.seed_mode;

      // 6) Merge with defaults, overlay, and user params (overlay wins over defaults, user wins over overlay)
      const requestParams = {
        ...defaultParams,
        ...(overlayParamsHttp || {}),
        ...filteredParams,
        seed: effectiveSeed,
        seed_mode: effectiveSeedMode,
      };

      // Capture effective model for audit
      effectiveModel = String((requestParams as any).model || "");

      // Extract LoRAs for audit: user > overlay > defaults
      try {
        const userSpecifiedLoras = Object.prototype.hasOwnProperty.call(
          filteredParams,
          "loras"
        );
        const overlaySelected = !!overlayParamsHttp;
        if (userSpecifiedLoras) {
          const ls = Array.isArray((filteredParams as any).loras)
            ? (filteredParams as any).loras
            : [];
          effectiveLoras = ls
            .map((x: any) => x?.file)
            .filter((x: any) => typeof x === "string" && x.trim())
            .map((x: any) => String(x).trim());
        } else if (overlaySelected) {
          const ls = Array.isArray((overlayParamsHttp as any)?.loras)
            ? (overlayParamsHttp as any).loras
            : [];
          effectiveLoras = ls
            .map((x: any) => x?.file)
            .filter((x: any) => typeof x === "string" && x.trim())
            .map((x: any) => String(x).trim());
        } else {
          // Fallback: extract LoRAs from requestParams (includes defaults)
          const ls = Array.isArray((requestParams as any)?.loras)
            ? (requestParams as any).loras
            : [];
          effectiveLoras = ls
            .map((x: any) => x?.file)
            .filter((x: any) => typeof x === "string" && x.trim())
            .map((x: any) => String(x).trim());
        }
      } catch {}

      console.debug(`use prompt: "${requestParams.prompt}"`);

      // 6) Call Draw Things API
      console.info("send request to Draw Things API...");
      startTime = Date.now();
      const response = await this.axios.post(
        "/sdapi/v1/txt2img",
        requestParams
      );
      endTime = Date.now();

      // 7) Validate response
      if (
        !response.data ||
        !response.data.images ||
        response.data.images.length === 0
      ) {
        throw new Error("API did not return image data");
      }

      const imagesArr: string[] = response.data.images as string[];
      console.info(
        `DrawThings: received ${imagesArr.length} image(s) from API`
      );

      // Normalize first image to buffer
      const first = imagesArr[0];
      const formattedFirst = first.startsWith("data:image/")
        ? first
        : `data:image/png;base64,${first}`;
      const cleanBase64 = formattedFirst.includes(",")
        ? formattedFirst.split(",")[1]
        : formattedFirst.replace(/^data:image\/\w+;base64,/, "");
      const imageBuffer = Buffer.from(cleanBase64, "base64");

      console.info("image generation success");

      // Audit: user input (raw form values)
      try {
        audit.setUserInput({
          width: (inputParams as any)?.width,
          height: (inputParams as any)?.height,
          imageFormat: (inputParams as any)?.imageFormat,
          quality: (inputParams as any)?.quality,
        });
      } catch {}

      // Audit: what we requested (payload-only, redacted)
      try {
        audit.setPrompt(String(requestParams.prompt || ""));
        audit.setRequested({
          source: "http.body",
          model: String((requestParams as any).model || ""),
          width: Number((requestParams as any).width),
          height: Number((requestParams as any).height),
          steps: Number((requestParams as any).steps),
        });
      } catch {}

      // Audit: what backend returned (image dims of first image)
      let backendW: number | undefined;
      let backendH: number | undefined;
      try {
        const meta = await imgGetSize(imageBuffer);
        backendW = meta.width;
        backendH = meta.height;
      } catch {}
      audit.setBackendReturned({ width: backendW, height: backendH });

      // Audit: what we post-processed (here identical to backend output)
      audit.setPostProcessed({ width: backendW, height: backendH });
      audit.setInferenceTime(endTime - startTime);

      // 8) Timestamps and metadata (measured above around HTTP call)

      const result: BackendGenerateResult = {
        isError: false,
        imageData: formattedFirst,
        imageBuffer,
        images: imagesArr.map((img) =>
          img.startsWith("data:image/") ? img : `data:image/png;base64,${img}`
        ),
        metadata: {
          alt: `Image generated from prompt: ${requestParams.prompt}`,
          inference_time_ms: endTime - startTime,
          model: effectiveModel,
          width: Number(requestParams.width),
          height: Number(requestParams.height),
          requested_dimensions: {
            width: Number(requestParams.width),
            height: Number(requestParams.height),
          },
          image_format: usedImageFormat,
          quality: usedQuality, // only set when mapping occurred; otherwise undefined
          steps: Number((requestParams as any).steps),
          seed:
            typeof (requestParams as any).seed === "number" &&
            Number.isFinite((requestParams as any).seed)
              ? (requestParams as any).seed
              : undefined,
          seed_mode: (() => {
            const sm = (requestParams as any).seed_mode ?? (requestParams as any).seedMode;
            return typeof sm === "string" && sm.trim() ? sm : undefined;
          })(),
          seed_source:
            source === "custom" &&
            typeof (overlayParamsHttp as any)?.seed === "number" &&
            Number.isFinite((overlayParamsHttp as any)?.seed)
              ? "custom"
              : "default",
          seed_mode_source:
            source === "custom" &&
            typeof (overlayParamsHttp as any)?.seed_mode === "string" &&
            String((overlayParamsHttp as any).seed_mode).trim()
              ? "custom"
              : "default",
          prompt_used: String(requestParams.prompt || ""),
          prompt_origin:
            typeof inputParams.prompt === "string" && inputParams.prompt.trim()
              ? "user"
              : "default:drawthings:txt2img",
          transport: "http",
          drawthings_params: requestParams,
          variantsReturned: imagesArr.length,
          // Overlay source tracking for audit
          overlay_source: overlaySource,
          ...(overlayPreset && { overlay_preset: overlayPreset }),
          defaults_used: "defaultParamsDrawThingsTxt2Img",
          overlay_lookup_mode: "txt2img",
          strength_used:
            typeof (requestParams as any).strength === "number" &&
            Number.isFinite((requestParams as any).strength)
              ? Number((requestParams as any).strength)
              : undefined,
          steps_used:
            typeof (requestParams as any).steps === "number" &&
            Number.isFinite((requestParams as any).steps)
              ? Number((requestParams as any).steps)
              : undefined,
          sampler_used:
            typeof (requestParams as any).sampler === "string" &&
            String((requestParams as any).sampler).trim()
              ? String((requestParams as any).sampler)
              : undefined,
          guidance_scale_used:
            typeof (requestParams as any).guidance_scale === "number" &&
            Number.isFinite((requestParams as any).guidance_scale)
              ? Number((requestParams as any).guidance_scale)
              : undefined,
          ...(effectiveLoras &&
            effectiveLoras.length > 0 && { loras_used: effectiveLoras }),
        },
      };
      // NOTE: Audit logging moved to core (tools.ts) to avoid duplication
      return result;
    } catch (error) {
      // Log only the message, not the full error object (avoids verbose stack traces in stderr)
      const errMsg = error instanceof Error ? error.message : String(error);
      console.error(`image generation error: ${errMsg}`);

      const axiosError = error as any;
      let status: number | undefined = undefined;
      let rawBody: string | undefined = undefined;

      if (axiosError && axiosError.response) {
        status = axiosError.response.status;
        const data = axiosError.response.data;
        if (typeof data === "string") rawBody = data;
        else if (data !== undefined) {
          try {
            rawBody = JSON.stringify(data);
          } catch {
            rawBody = String(data);
          }
        }
      }

      if (!rawBody) {
        if (typeof axiosError?.message === "string")
          rawBody = axiosError.message;
        else if (error instanceof Error && typeof error.message === "string")
          rawBody = error.message;
        else rawBody = "unknown error";
      }

      // NOTE: Error audit logging moved to core (tools.ts) to avoid duplication
      return {
        isError: true,
        status,
        errorMessage: rawBody,
      };
    }
  }

  /**
   * generate image (image-to-image)
   * @param inputParams user provided params (raw from tool).
   * @param sourceBuffer PNG/JPG/WebP image as Buffer; will be converted to PNG and sent as base64
   * @param _onProgress Progress callback (not supported by HTTP backend, ignored)
   */
  async generateImageImg2Img(
    inputParams: Partial<ImageGenerationParams> = {},
    sourceBuffer: Buffer,
    _onProgress?: ProgressCallback
  ): Promise<BackendGenerateResult> {
    const audit = buildAuditLogger({ backend: this.name, mode: "img2img" });
    // Overlay tracking for result metadata
    let overlaySource: "custom" | "modelOverlay" | "default" = "default";
    let overlayPreset: string | undefined;
    let effectiveModel: string | undefined;
    let effectiveLoras: string[] | undefined;

    try {
      const i2iProfile = inputParams?._dt_i2i_profile;
      const baseDefaults =
        i2iProfile === "edit" ? defaultParamsEdit : defaultParamsImg2Img;
      const baseLimits =
        i2iProfile === "edit" ? drawthingsEditLimits : drawthingsLimits;

      // Measure only backend latency (request -> response), like gRPC path
      let startTime = 0;
      let endTime = 0;
      // Validate + mapping (imageFormat, quality, variants)
      let params: Partial<ImageGenerationParams> = {};
      try {
        const validationResult = validateImageGenerationParams(inputParams);
        if (validationResult.valid) {
          params = inputParams;
        }
      } catch {}

      if (
        params.random_string &&
        (!params.prompt || Object.keys(params).length === 1)
      ) {
        params.prompt = params.random_string;
        delete params.random_string;
      }
      if (!params.prompt) {
        params.prompt = inputParams.prompt || baseDefaults.prompt;
      }

      let usedImageFormat: string | undefined = undefined;
      let usedQuality: string | undefined = undefined;
      try {
        const fmt = (inputParams as any).imageFormat as string | undefined;
        // Shorthand should not override explicit dimensions (core passes requested_effective via width/height).
        if ((params as any).width == null && (params as any).height == null) {
          if (fmt === "square") {
            params.width = 1024;
            params.height = 1024;
            usedImageFormat = "square";
          } else if (fmt === "landscape") {
            params.width = 1024;
            params.height = 768;
            usedImageFormat = "landscape";
          } else if (fmt === "portrait") {
            params.width = 768;
            params.height = 1024;
            usedImageFormat = "portrait";
          } else if (fmt === "16:9") {
            params.width = 1024;
            params.height = 576;
            usedImageFormat = "16:9";
          }
        }

        const qual = (inputParams as any).quality as string | undefined;
        if (qual === "low") {
          params.steps = 4;
          usedQuality = "low";
        } else if (qual === "medium") {
          params.steps = 8;
          usedQuality = "medium";
        } else if (qual === "high") {
          params.steps = 12;
          usedQuality = "high";
        }

        const variants = (inputParams as any).variants as number | undefined;
        if (typeof variants === "number") {
          const v = Math.max(1, Math.min(4, Math.round(variants)));
          (params as any).batch_size = v;
          (params as any).batch_count = 1;
        }
      } catch {}

      // Filter to allowed keys from defaults
      const allowedKeys = new Set(Object.keys(baseDefaults));
      const filteredParams: Partial<ImageGenerationParams> = {};
      for (const [k, v] of Object.entries(params)) {
        if (allowedKeys.has(k)) (filteredParams as any)[k] = v as any;
      }

      // Apply effective overlay (Custom Configs → Model Overlay → Defaults)
      // Use i2iProfile to determine correct mode for overlay lookup
      const modelId = (inputParams as any).model as ModelId | undefined;
      const overlayMode = i2iProfile === "edit" ? "edit" : "img2img";
      const {
        source,
        presetName,
        params: overlayParams,
      } = getEffectiveOverlay(modelId, overlayMode);

      const overlayParamsHttp =
        source === "custom" && overlayParams
          ? applyHttpCustomOverlayFixups(overlayParams)
          : overlayParams;

      // Capture overlay info for result metadata
      overlaySource = source;
      overlayPreset = presetName;

      // Tool-level preset 'auto' means: do not override engine model.
      if (modelId === "auto") {
        delete (filteredParams as any).model;
      }
      if (overlayParamsHttp) {
        console.info(
          `Applying ${source} overlay for '${modelId ?? 'auto'}' (img2img)${
            presetName ? ` preset=${presetName}` : ""
          }`
        );
        // Remove 'model' from filteredParams so overlay model filename wins
        delete (filteredParams as any).model;
      }

      // Log overlay source to audit
      audit.setExtras({
        overlay_source: source,
        ...(presetName && { overlay_preset: presetName }),
      });

      // Ensure PNG base64 for init_images
      let initPng = sourceBuffer;
      let srcW: number | undefined;
      let srcH: number | undefined;
      try {
        // Derive dimensions from the actual PNG we send.
        initPng = await imgToPng(sourceBuffer);
        const meta = await imgGetSize(initPng);
        srcW = meta.width;
        srcH = meta.height;
      } catch {}
      const initB64 = initPng.toString("base64");

      // Merge payload (defaults → overlay → user params)
      const requestParams: any = {
        ...baseDefaults,
        ...(overlayParamsHttp || {}),
        ...filteredParams,
        // Seed policy: allow Custom Configs to override (incl. -1); otherwise keep defaults.
        seed:
          source === "custom" &&
          typeof (overlayParamsHttp as any)?.seed === "number" &&
          Number.isFinite((overlayParamsHttp as any)?.seed)
            ? (overlayParamsHttp as any).seed
            : baseDefaults.seed,
        // Seed mode policy mirrors seed: allow Custom Configs to override; otherwise keep defaults.
        seed_mode:
          source === "custom" &&
          typeof (overlayParamsHttp as any)?.seed_mode === "string" &&
          String((overlayParamsHttp as any).seed_mode).trim()
            ? String((overlayParamsHttp as any).seed_mode)
            : (baseDefaults as any).seed_mode,
        init_images: [initB64],
      };

      const defaultsUsed =
        i2iProfile === "edit"
          ? "defaultParamsDrawThingsEdit"
          : "defaultParamsDrawThingsImg2Img";
      const overlayLookupMode = overlayMode;

      // Capture effective model for audit
      effectiveModel = String(requestParams.model || "");

      // Extract LoRAs for audit: user > overlay > defaults
      try {
        const userSpecifiedLoras = Object.prototype.hasOwnProperty.call(
          filteredParams,
          "loras"
        );
        const overlaySelected = !!overlayParamsHttp;
        if (userSpecifiedLoras) {
          const ls = Array.isArray((filteredParams as any).loras)
            ? (filteredParams as any).loras
            : [];
          effectiveLoras = ls
            .map((x: any) => x?.file)
            .filter((x: any) => typeof x === "string" && x.trim())
            .map((x: any) => String(x).trim());
        } else if (overlaySelected) {
          const ls = Array.isArray((overlayParamsHttp as any)?.loras)
            ? (overlayParamsHttp as any).loras
            : [];
          effectiveLoras = ls
            .map((x: any) => x?.file)
            .filter((x: any) => typeof x === "string" && x.trim())
            .map((x: any) => String(x).trim());
        } else {
          // Fallback: extract LoRAs from requestParams (includes defaults)
          const ls = Array.isArray((requestParams as any)?.loras)
            ? (requestParams as any).loras
            : [];
          effectiveLoras = ls
            .map((x: any) => x?.file)
            .filter((x: any) => typeof x === "string" && x.trim())
            .map((x: any) => String(x).trim());
        }
      } catch {}

      // Respect explicit requested output size when provided; otherwise mirror source
      const userReqW = (() => {
        const v = (filteredParams as any)?.width ?? (params as any)?.width;
        return typeof v === "number" && Number.isFinite(v)
          ? Math.round(v)
          : undefined;
      })();
      const userReqH = (() => {
        const v = (filteredParams as any)?.height ?? (params as any)?.height;
        return typeof v === "number" && Number.isFinite(v)
          ? Math.round(v)
          : undefined;
      })();

      // Dimension contract (HTTP):
      // Draw Things HTTP rejects img2img requests when width/height don't match init_images.
      // We therefore force width/height/target_* to match the init image exactly.
      const hasInitImages =
        Array.isArray(requestParams.init_images) &&
        requestParams.init_images.length > 0;
      if (srcW && srcH) {
        requestParams.original_width = srcW;
        requestParams.original_height = srcH;
      }
      if (hasInitImages) {
        if (srcW && srcH) {
          requestParams.width = srcW;
          requestParams.height = srcH;
          requestParams.target_width = srcW;
          requestParams.target_height = srcH;
        }
      } else if (userReqW && userReqH) {
        requestParams.width = userReqW;
        requestParams.height = userReqH;
        requestParams.target_width = userReqW;
        requestParams.target_height = userReqH;
      }

      // Upscaler is controlled deterministically by core.
      // No heuristics/fallbacks here to avoid legacy drift.
      {
        const needs = (inputParams as any)?._dt_needs_upscaler;
        if (typeof needs !== "boolean") {
          throw new Error(
            "Invariant failed: _dt_needs_upscaler must be provided by core for img2img"
          );
        }
        if (needs) {
          requestParams.upscaler = baseLimits.upscaler;
          requestParams.upscaler_scale = baseLimits.upscalerScaleFactor;
          console.info(`DT img2img upscaler enabled (core)`);
        } else {
          requestParams.upscaler = null;
          requestParams.upscaler_scale = 0;
          console.info(`DT img2img upscaler disabled (core)`);
        }
      }

      console.info(
        `DT img2img size set: width=${requestParams.width}, height=${
          requestParams.height
        }, original_width=${requestParams.original_width}, original_height=${
          requestParams.original_height
        }, target_width=${requestParams.target_width}, target_height=${
          requestParams.target_height
        }, upscaler=${requestParams.upscaler ?? "-"}, upscaler_scale=${
          requestParams.upscaler_scale ?? "-"
        }`
      );
      console.info("send request to Draw Things API (img2img)...");
      startTime = Date.now();
      const response = await this.axios.post(
        "/sdapi/v1/img2img",
        requestParams
      );
      endTime = Date.now();

      if (
        !response.data ||
        !response.data.images ||
        response.data.images.length === 0
      ) {
        throw new Error("API did not return image data");
      }

      const imagesArr: string[] = response.data.images as string[];
      const first = imagesArr[0];
      const formattedFirst = first.startsWith("data:image/")
        ? first
        : `data:image/png;base64,${first}`;
      const cleanBase64 = formattedFirst.includes(",")
        ? formattedFirst.split(",")[1]
        : formattedFirst.replace(/^data:image\/\w+;base64,/, "");
      const imageBuffer = Buffer.from(cleanBase64, "base64");

      // measured above around HTTP call

      // Audit: requested payload summary
      try {
        audit.setPrompt(String(requestParams.prompt || ""));
        audit.setRequested({
          source: "http.body",
          model: String((requestParams as any).model || ""),
          width: Number((requestParams as any).width),
          height: Number((requestParams as any).height),
          steps: Number((requestParams as any).steps),
          original_width: Number((requestParams as any).original_width),
          original_height: Number((requestParams as any).original_height),
          target_width: Number((requestParams as any).target_width),
          target_height: Number((requestParams as any).target_height),
        });
        if (srcW && srcH) {
          audit.setExtras({ srcWidth: srcW, srcHeight: srcH });
        }
      } catch {}

      // Audit: backend returned dims
      let backendW: number | undefined;
      let backendH: number | undefined;
      try {
        const meta = await imgGetSize(imageBuffer);
        backendW = meta.width;
        backendH = meta.height;
      } catch {}
      audit.setBackendReturned({ width: backendW, height: backendH });

      // Post-processed: here identical to backend output (no resize applied here)
      audit.setPostProcessed({ width: backendW, height: backendH });
      audit.setInferenceTime(endTime - startTime);

      const result: BackendGenerateResult = {
        isError: false,
        imageData: formattedFirst,
        imageBuffer,
        images: imagesArr.map((img) =>
          img.startsWith("data:image/") ? img : `data:image/png;base64,${img}`
        ),
        metadata: {
          alt: `Image-to-image from prompt: ${requestParams.prompt}`,
          inference_time_ms: endTime - startTime,
          model: effectiveModel,
          width: Number(requestParams.width),
          height: Number(requestParams.height),
          requested_dimensions: {
            width: Number(requestParams.width),
            height: Number(requestParams.height),
          },
          image_format: usedImageFormat,
          quality: usedQuality, // only set when mapping occurred; otherwise undefined
          steps: Number((requestParams as any).steps),
          seed:
            typeof (requestParams as any).seed === "number" &&
            Number.isFinite((requestParams as any).seed)
              ? (requestParams as any).seed
              : undefined,
          seed_mode: (() => {
            const sm = (requestParams as any).seed_mode ?? (requestParams as any).seedMode;
            return typeof sm === "string" && sm.trim() ? sm : undefined;
          })(),
          prompt_used: String(requestParams.prompt || ""),
          prompt_origin:
            typeof inputParams.prompt === "string" && inputParams.prompt.trim()
              ? "user"
              : "default:drawthings:img2img",
          transport: "http",
          // Overlay source tracking for audit
          overlay_source: overlaySource,
          ...(overlayPreset && { overlay_preset: overlayPreset }),
          defaults_used: defaultsUsed,
          overlay_lookup_mode: overlayLookupMode,
          i2i_profile: i2iProfile,
          strength_used:
            typeof requestParams.strength === "number" &&
            Number.isFinite(requestParams.strength)
              ? requestParams.strength
              : undefined,
          steps_used:
            typeof requestParams.steps === "number" &&
            Number.isFinite(requestParams.steps)
              ? requestParams.steps
              : undefined,
          sampler_used:
            typeof requestParams.sampler === "string" &&
            String(requestParams.sampler).trim()
              ? String(requestParams.sampler)
              : undefined,
          guidance_scale_used:
            typeof requestParams.guidance_scale === "number" &&
            Number.isFinite(requestParams.guidance_scale)
              ? requestParams.guidance_scale
              : undefined,
          ...(effectiveLoras &&
            effectiveLoras.length > 0 && { loras_used: effectiveLoras }),
        },
      };
      // Do NOT write audit here; index.ts will write a consolidated audit entry
      return result;
    } catch (error) {
      // Log only the message, not the full error object (avoids verbose stack traces in stderr)
      const errMsg = error instanceof Error ? error.message : String(error);
      console.error(`image2image generation error: ${errMsg}`);
      const axiosError = error as any;
      let status: number | undefined = undefined;
      let rawBody: string | undefined = undefined;

      if (axiosError && axiosError.response) {
        status = axiosError.response.status;
        const data = axiosError.response.data;
        if (typeof data === "string") rawBody = data;
        else if (data !== undefined) {
          try {
            rawBody = JSON.stringify(data);
          } catch {
            rawBody = String(data);
          }
        }
      }

      if (!rawBody) {
        if (typeof axiosError?.message === "string")
          rawBody = axiosError.message;
        else if (error instanceof Error && typeof error.message === "string")
          rawBody = error.message;
        else rawBody = "unknown error";
      }

      // Skip audit write here as well; index.ts handles consolidated logging
      return { isError: true, status, errorMessage: rawBody };
    }
  }
}
export default DrawThingsService;