import { tool, type Tool, type ToolsProviderController, type ToolCallContext } from "@lmstudio/sdk";
import { z } from "zod";
import {
  getBackgroundRequest,
  markBackgroundRequestDelivered,
  listBackgroundRequests,
  cancelBackgroundRequest,
  requestBelongsToScope,
  formatToolMetaBlock,
} from "../core-bundle.mjs";
import path from "path";

const COMPLETED_LIST_HINT =
  "This list shows statuses only. To evaluate a finished (completed) image, call request_progress with its requestId.";

export function createRequestProgressTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "request_progress",
    description: `Check generate_image background request(s) started with runInBackground: true.

- Normal use: pass requestId (the id generate_image returned). While the render runs this shows its progress; once it has finished it returns the full result including the image.
- To CANCEL a request (queued or running): pass requestId and cancel: true.
- Without requestId the tool only lists the status of all requests for this chat (no images); use that just to recover a lost requestId.

Do not call this again immediately after a non-completed result — wait for the user's next message or other work first.

${formatToolMetaBlock()}`,
    parameters: {
      requestId: z.string().optional().describe("The requestId returned by generate_image. Omit to list every background request for this chat."),
      cancel: z.boolean().optional().describe("If true, cancels the request given by requestId instead of checking it. Queued requests are removed from the queue, running ones are aborted."),
    },
    implementation: async (args: any, _ctx: ToolCallContext) => {
      try {
        const workingDir = ctl.getWorkingDirectory();
        if (typeof workingDir !== "string" || !workingDir.trim()) {
          return "request_progress failed: no active chat working directory.";
        }
        const requestId = String(args?.requestId ?? "").trim();

        if (!requestId && args?.cancel === true) {
          return "request_progress failed: cancel requires a requestId.";
        }

        if (!requestId) {
          const all = await listBackgroundRequests(workingDir);
          if (all.length === 0) return "No background requests found for this chat.";
          const lines = all.map((r) => `- ${r.requestId} [${r.status === "running" && r.cancelRequested ? "cancelling" : r.status}]${r.status === "running" && r.message ? `: ${r.message}` : ""}`).join("\n");
          return all.some((r) => r.status === "completed") ? `${lines}\n${COMPLETED_LIST_HINT}` : lines;
        }

        const chatId = path.basename(workingDir);
        if (!requestBelongsToScope(requestId, chatId)) {
          return `requestId "${requestId}" does not belong to this chat (${chatId}).`;
        }

        if (args?.cancel === true) {
          const res = await cancelBackgroundRequest(workingDir, requestId);
          if (res.outcome === "not-found") {
            return `No background request found for requestId "${requestId}". It may have expired or never existed.`;
          }
          if (res.outcome === "already-finished") {
            return `generate_image (requestId ${requestId}) is already ${res.status}; nothing to cancel.`;
          }
          if (res.outcome === "cancelled-queued") {
            return `generate_image (requestId ${requestId}) was removed from the queue and cancelled.`;
          }
          if (res.outcome === "cancelled") {
            return `generate_image (requestId ${requestId}) was cancelled: the job was aborted and has ended.`;
          }
          if (res.outcome === "finished-before-cancel") {
            return `generate_image (requestId ${requestId}) had already ${res.status} before the cancel took effect; nothing was cancelled. Use request_progress with this requestId to retrieve it.`;
          }
          return `Cancellation of generate_image (requestId ${requestId}) was requested, but the job was still running after ${Math.round(res.waitedMs / 1000)} s; it is not confirmed as cancelled. Check again with request_progress.`;
        }

        const record = await getBackgroundRequest(workingDir, requestId);
        if (!record) {
          return `No background request found for requestId "${requestId}". It may have expired or never existed.`;
        }

        if (record.status === "cancelled") {
          return `generate_image (requestId ${requestId}) was cancelled.`;
        }

        if (record.status === "completed") {
          await markBackgroundRequestDelivered(workingDir, requestId);
          const resp = record.result as any;
          if (resp && Array.isArray(resp.content)) return resp.content;
          return typeof resp === "string" ? resp : JSON.stringify(resp);
        }

        if (record.status === "failed") {
          await markBackgroundRequestDelivered(workingDir, requestId);
          return `generate_image (requestId ${requestId}) failed: ${record.error || "unknown error"}`;
        }

        if (record.status === "running") {
          const stepInfo =
            record.message ||
            (typeof record.step === "number" && typeof record.totalSteps === "number" && record.totalSteps > 0
              ? `step ${record.step}/${record.totalSteps}`
              : "rendering");
          return `generate_image (requestId ${requestId}) is still running: ${stepInfo}. Do not call request_progress again immediately — wait for the user's next message or other work before checking back.`;
        }

        return `generate_image (requestId ${requestId}) is queued and hasn't started rendering yet. Do not call request_progress again immediately — wait for the user's next message or other work before checking back.`;
      } catch (e) {
        return `request_progress failed: ${String((e as any)?.message || e)}`;
      }
    },
  });
}

import { tool, type Tool, type ToolsProviderController, type ToolCallContext } from "@lmstudio/sdk";
import { z } from "zod";
import {
  getBackgroundRequest,
  markBackgroundRequestDelivered,
  listBackgroundRequests,
  cancelBackgroundRequest,
  requestBelongsToScope,
  formatToolMetaBlock,
} from "../core-bundle.mjs";
import path from "path";

const COMPLETED_LIST_HINT =
  "This list shows statuses only. To evaluate a finished (completed) image, call request_progress with its requestId.";

export function createRequestProgressTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "request_progress",
    description: `Check generate_image background request(s) started with runInBackground: true.

- Normal use: pass requestId (the id generate_image returned). While the render runs this shows its progress; once it has finished it returns the full result including the image.
- To CANCEL a request (queued or running): pass requestId and cancel: true.
- Without requestId the tool only lists the status of all requests for this chat (no images); use that just to recover a lost requestId.

Do not call this again immediately after a non-completed result — wait for the user's next message or other work first.

${formatToolMetaBlock()}`,
    parameters: {
      requestId: z.string().optional().describe("The requestId returned by generate_image. Omit to list every background request for this chat."),
      cancel: z.boolean().optional().describe("If true, cancels the request given by requestId instead of checking it. Queued requests are removed from the queue, running ones are aborted."),
    },
    implementation: async (args: any, _ctx: ToolCallContext) => {
      try {
        const workingDir = ctl.getWorkingDirectory();
        if (typeof workingDir !== "string" || !workingDir.trim()) {
          return "request_progress failed: no active chat working directory.";
        }
        const requestId = String(args?.requestId ?? "").trim();

        if (!requestId && args?.cancel === true) {
          return "request_progress failed: cancel requires a requestId.";
        }

        if (!requestId) {
          const all = await listBackgroundRequests(workingDir);
          if (all.length === 0) return "No background requests found for this chat.";
          const lines = all.map((r) => `- ${r.requestId} [${r.status === "running" && r.cancelRequested ? "cancelling" : r.status}]${r.status === "running" && r.message ? `: ${r.message}` : ""}`).join("\n");
          return all.some((r) => r.status === "completed") ? `${lines}\n${COMPLETED_LIST_HINT}` : lines;
        }

        const chatId = path.basename(workingDir);
        if (!requestBelongsToScope(requestId, chatId)) {
          return `requestId "${requestId}" does not belong to this chat (${chatId}).`;
        }

        if (args?.cancel === true) {
          const res = await cancelBackgroundRequest(workingDir, requestId);
          if (res.outcome === "not-found") {
            return `No background request found for requestId "${requestId}". It may have expired or never existed.`;
          }
          if (res.outcome === "already-finished") {
            return `generate_image (requestId ${requestId}) is already ${res.status}; nothing to cancel.`;
          }
          if (res.outcome === "cancelled-queued") {
            return `generate_image (requestId ${requestId}) was removed from the queue and cancelled.`;
          }
          if (res.outcome === "cancelled") {
            return `generate_image (requestId ${requestId}) was cancelled: the job was aborted and has ended.`;
          }
          if (res.outcome === "finished-before-cancel") {
            return `generate_image (requestId ${requestId}) had already ${res.status} before the cancel took effect; nothing was cancelled. Use request_progress with this requestId to retrieve it.`;
          }
          return `Cancellation of generate_image (requestId ${requestId}) was requested, but the job was still running after ${Math.round(res.waitedMs / 1000)} s; it is not confirmed as cancelled. Check again with request_progress.`;
        }

        const record = await getBackgroundRequest(workingDir, requestId);
        if (!record) {
          return `No background request found for requestId "${requestId}". It may have expired or never existed.`;
        }

        if (record.status === "cancelled") {
          return `generate_image (requestId ${requestId}) was cancelled.`;
        }

        if (record.status === "completed") {
          await markBackgroundRequestDelivered(workingDir, requestId);
          const resp = record.result as any;
          if (resp && Array.isArray(resp.content)) return resp.content;
          return typeof resp === "string" ? resp : JSON.stringify(resp);
        }

        if (record.status === "failed") {
          await markBackgroundRequestDelivered(workingDir, requestId);
          return `generate_image (requestId ${requestId}) failed: ${record.error || "unknown error"}`;
        }

        if (record.status === "running") {
          const stepInfo =
            record.message ||
            (typeof record.step === "number" && typeof record.totalSteps === "number" && record.totalSteps > 0
              ? `step ${record.step}/${record.totalSteps}`
              : "rendering");
          return `generate_image (requestId ${requestId}) is still running: ${stepInfo}. Do not call request_progress again immediately — wait for the user's next message or other work before checking back.`;
        }

        return `generate_image (requestId ${requestId}) is queued and hasn't started rendering yet. Do not call request_progress again immediately — wait for the user's next message or other work before checking back.`;
      } catch (e) {
        return `request_progress failed: ${String((e as any)?.message || e)}`;
      }
    },
  });
}