import { tool, type Tool, type ToolsProviderController, type ToolCallContext } from "@lmstudio/sdk";
import { z } from "zod";
import {
  GenerateToolParamsShapeMinimal,
  setActiveChatContext,
  formatToolMetaBlock,
  syncAttachmentsToState,
  createBackgroundRequestId,
  runInBackgroundJob,
  createStatusProgressCallback,
} from "../core-bundle.mjs";
import path from "path";

// Temporarily hide `quality` from the agent-facing interface.
// Internally, the backend + handler still support it.
// eslint-disable-next-line @typescript-eslint/no-unused-vars
const { quality: _hiddenQuality, ...GenerateToolParamsShapeAgent } =
  GenerateToolParamsShapeMinimal;

export function createGenerateImageTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "generate_image",
    description: `Generate an image or video using Draw Things.

All parameters have sensible defaults. Only override them when there's a good reason OR the user explicitly requests it.

Parameters:
- prompt: Image/video description (mode: text2image/text2video) OR description of desired changes (mode: image2image/edit/image2video).
- mode:
  - text2image: generate a NEW image from the prompt. No sources needed.
  - image2image: uses source image(s) as visual context. Supports moodboard via gRPC if model allows. Prompting: focus on desired refinement.
  - edit: Supports MULTIPLE reference images (canvas + moodboard). Prompting: apply only the requested changes, be short, concrete, action-oriented.
  - text2video: generate a video from the prompt. No sources needed. ltx: 25 fps. minimax/minimax-ref: 24 fps + 32 kHz audio.
  - image2video: generate a video from a source image + prompt. ltx: single canvas image, 25 fps. minimax: canvas = first frame, first enabled moodboard image = last frame (2 reference images max), 24 fps + 32 kHz audio. minimax-ref: subject/object/scene references via canvas + moodboard (4 reference images max), 24 fps + 32 kHz audio.
- model: auto (default) | z-image | krea | qwen-image | qwen-image-2.1 | qwen-image-2.1-turbo | flux | ltx | minimax | minimax-ref | custom | <user-defined>. auto selects best model for the mode. User-defined entries (e.g. portrait or anime) map to user-defined models in custom_configs.json.
- imageFormat: square | landscape | portrait | 16:9. Override if context suggests. 16:9 yields 1024×576 (video-optimized).
- variants: 1-4. Number of images. Default is 1. Note: text2video/image2video support only 1 variant; higher values are silently ignored by the backend.
- numFrames: Number of video frames (text2video/image2video only). Must be ×32+1 (e.g. 33, 65, 97). Invalid values are auto-corrected.
- canvas: Primary source image for image2image/edit/image2video. REQUIRED for these modes.
  - Notation: a1 (attachment 1), v2 (variant 2), p1 (picture 1), i3 (image 3)
  - Digit-only (e.g., 2) allowed only when a single pool is populated
- moodboard: Additional style references for image2image/edit modes. Array of same notations as canvas.
- runInBackground: If true, returns immediately with a requestId instead of waiting for the render to finish. Poll with request_progress(requestId).

Critical rules:
- HTTP transport: image2image supports only ONE source. Use gRPC or mode=edit for multiple references.
- Always specify canvas explicitly for image2image/edit/image2video modes.

Returns: inline preview(s), JSON with file URLs, and clickable links.

${formatToolMetaBlock()}`,
    parameters: GenerateToolParamsShapeAgent as Record<string, z.ZodTypeAny>,
    implementation: async (args: any, ctx: ToolCallContext) => {
      try {
        // CRITICAL: Refresh chat context on every tool call to prevent TTL expiry.
        // The Generator sets context once per turn, but tool execution can take longer
        // than the 60s TTL (model thinking + Draw Things generation time).
        const workingDir = ctl.getWorkingDirectory();
        try {
          if (typeof workingDir === "string" && workingDir.trim().length > 0) {
            const chatId = path.basename(workingDir);
            if (/^\d+$/.test(chatId)) {
              setActiveChatContext({
                chatId,
                workingDir,
                requestId: `tool-${Date.now()}`,
              });
            }
          }
        } catch {
          // non-fatal: fallback to existing context or heuristic
        }

        // Sync attachments from conversation.json → chat_media_state.json.
        // Non-fatal: a sync failure must not block image generation.
        try {
          if (typeof workingDir === "string" && workingDir.trim().length > 0) {
            await syncAttachmentsToState(workingDir, false, Number.MAX_SAFE_INTEGER);
          }
        } catch (syncErr: any) {
          console.warn(
            "[generate_image] attachment sync failed (non-fatal):",
            syncErr?.message ?? syncErr,
          );
        }

        // Create progress callback that updates the tool status in the UI.
        // Note: Only gRPC backend invokes this callback (streaming).
        // HTTP backend has no progress streaming, so it stays silent.
        const onProgress = createStatusProgressCallback(ctx);

        const mod = await import("../core/tools.js");
        if (typeof (mod as any).handleGenerateImage !== "function") {
          return "generate_image core handler not available.";
        }

        if (args?.runInBackground === true) {
          if (typeof workingDir !== "string" || !workingDir.trim()) {
            return "generate_image failed: runInBackground requires an active chat working directory.";
          }
          const requestId = createBackgroundRequestId(path.basename(workingDir));
          // Awaited: guarantees the record exists on disk before this tool call returns, so an
          // immediate request_progress call in the same turn can never race ahead of it.
          await runInBackgroundJob({
            chatWd: workingDir,
            requestId,
            tool: "generate_image",
            run: (bgOnProgress) =>
              (mod as any).handleGenerateImage(args, bgOnProgress, { requestIdOverride: requestId }),
          });
          return (
            `Generation started in the background (requestId: ${requestId}). ` +
            `Use request_progress with this requestId to check status or retrieve the result once it's done. ` +
            `Do not call request_progress right away — wait for the user's next message or do other work first.`
          );
        }

        const resp = await (mod as any).handleGenerateImage(args, onProgress);
        // Return structured content so LM Studio can auto-save images and render markdown
        if (resp && Array.isArray((resp as any).content))
          return (resp as any).content;
        return typeof resp === "string" ? resp : JSON.stringify(resp);
      } catch (e) {
        return `generate_image failed: ${String((e as any)?.message || e)}`;
      }
    },
  });
}

import { tool, type Tool, type ToolsProviderController, type ToolCallContext } from "@lmstudio/sdk";
import { z } from "zod";
import {
  GenerateToolParamsShapeMinimal,
  setActiveChatContext,
  formatToolMetaBlock,
  syncAttachmentsToState,
  createBackgroundRequestId,
  runInBackgroundJob,
  createStatusProgressCallback,
} from "../core-bundle.mjs";
import path from "path";

// Temporarily hide `quality` from the agent-facing interface.
// Internally, the backend + handler still support it.
// eslint-disable-next-line @typescript-eslint/no-unused-vars
const { quality: _hiddenQuality, ...GenerateToolParamsShapeAgent } =
  GenerateToolParamsShapeMinimal;

export function createGenerateImageTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "generate_image",
    description: `Generate an image or video using Draw Things.

All parameters have sensible defaults. Only override them when there's a good reason OR the user explicitly requests it.

Parameters:
- prompt: Image/video description (mode: text2image/text2video) OR description of desired changes (mode: image2image/edit/image2video).
- mode:
  - text2image: generate a NEW image from the prompt. No sources needed.
  - image2image: uses source image(s) as visual context. Supports moodboard via gRPC if model allows. Prompting: focus on desired refinement.
  - edit: Supports MULTIPLE reference images (canvas + moodboard). Prompting: apply only the requested changes, be short, concrete, action-oriented.
  - text2video: generate a video from the prompt. No sources needed. ltx: 25 fps. minimax/minimax-ref: 24 fps + 32 kHz audio.
  - image2video: generate a video from a source image + prompt. ltx: single canvas image, 25 fps. minimax: canvas = first frame, first enabled moodboard image = last frame (2 reference images max), 24 fps + 32 kHz audio. minimax-ref: subject/object/scene references via canvas + moodboard (4 reference images max), 24 fps + 32 kHz audio.
- model: auto (default) | z-image | krea | qwen-image | qwen-image-2.1 | qwen-image-2.1-turbo | flux | ltx | minimax | minimax-ref | custom | <user-defined>. auto selects best model for the mode. User-defined entries (e.g. portrait or anime) map to user-defined models in custom_configs.json.
- imageFormat: square | landscape | portrait | 16:9. Override if context suggests. 16:9 yields 1024×576 (video-optimized).
- variants: 1-4. Number of images. Default is 1. Note: text2video/image2video support only 1 variant; higher values are silently ignored by the backend.
- numFrames: Number of video frames (text2video/image2video only). Must be ×32+1 (e.g. 33, 65, 97). Invalid values are auto-corrected.
- canvas: Primary source image for image2image/edit/image2video. REQUIRED for these modes.
  - Notation: a1 (attachment 1), v2 (variant 2), p1 (picture 1), i3 (image 3)
  - Digit-only (e.g., 2) allowed only when a single pool is populated
- moodboard: Additional style references for image2image/edit modes. Array of same notations as canvas.
- runInBackground: If true, returns immediately with a requestId instead of waiting for the render to finish. Poll with request_progress(requestId).

Critical rules:
- HTTP transport: image2image supports only ONE source. Use gRPC or mode=edit for multiple references.
- Always specify canvas explicitly for image2image/edit/image2video modes.

Returns: inline preview(s), JSON with file URLs, and clickable links.

${formatToolMetaBlock()}`,
    parameters: GenerateToolParamsShapeAgent as Record<string, z.ZodTypeAny>,
    implementation: async (args: any, ctx: ToolCallContext) => {
      try {
        // CRITICAL: Refresh chat context on every tool call to prevent TTL expiry.
        // The Generator sets context once per turn, but tool execution can take longer
        // than the 60s TTL (model thinking + Draw Things generation time).
        const workingDir = ctl.getWorkingDirectory();
        try {
          if (typeof workingDir === "string" && workingDir.trim().length > 0) {
            const chatId = path.basename(workingDir);
            if (/^\d+$/.test(chatId)) {
              setActiveChatContext({
                chatId,
                workingDir,
                requestId: `tool-${Date.now()}`,
              });
            }
          }
        } catch {
          // non-fatal: fallback to existing context or heuristic
        }

        // Sync attachments from conversation.json → chat_media_state.json.
        // Non-fatal: a sync failure must not block image generation.
        try {
          if (typeof workingDir === "string" && workingDir.trim().length > 0) {
            await syncAttachmentsToState(workingDir, false, Number.MAX_SAFE_INTEGER);
          }
        } catch (syncErr: any) {
          console.warn(
            "[generate_image] attachment sync failed (non-fatal):",
            syncErr?.message ?? syncErr,
          );
        }

        // Create progress callback that updates the tool status in the UI.
        // Note: Only gRPC backend invokes this callback (streaming).
        // HTTP backend has no progress streaming, so it stays silent.
        const onProgress = createStatusProgressCallback(ctx);

        const mod = await import("../core/tools.js");
        if (typeof (mod as any).handleGenerateImage !== "function") {
          return "generate_image core handler not available.";
        }

        if (args?.runInBackground === true) {
          if (typeof workingDir !== "string" || !workingDir.trim()) {
            return "generate_image failed: runInBackground requires an active chat working directory.";
          }
          const requestId = createBackgroundRequestId(path.basename(workingDir));
          // Awaited: guarantees the record exists on disk before this tool call returns, so an
          // immediate request_progress call in the same turn can never race ahead of it.
          await runInBackgroundJob({
            chatWd: workingDir,
            requestId,
            tool: "generate_image",
            run: (bgOnProgress) =>
              (mod as any).handleGenerateImage(args, bgOnProgress, { requestIdOverride: requestId }),
          });
          return (
            `Generation started in the background (requestId: ${requestId}). ` +
            `Use request_progress with this requestId to check status or retrieve the result once it's done. ` +
            `Do not call request_progress right away — wait for the user's next message or do other work first.`
          );
        }

        const resp = await (mod as any).handleGenerateImage(args, onProgress);
        // Return structured content so LM Studio can auto-save images and render markdown
        if (resp && Array.isArray((resp as any).content))
          return (resp as any).content;
        return typeof resp === "string" ? resp : JSON.stringify(resp);
      } catch (e) {
        return `generate_image failed: ${String((e as any)?.message || e)}`;
      }
    },
  });
}