Peekaboo/tests/unit/tools/image.test.ts

import { vi } from "vitest";
import {
  imageToolHandler,
  buildSwiftCliArgs,
  ImageToolInput,
} from "../../../src/tools/image";
import {
  executeSwiftCli,
  readImageAsBase64,
} from "../../../src/utils/peekaboo-cli";
import { mockSwiftCli } from "../../mocks/peekaboo-cli.mock";
import { pino } from "pino";
import {
  SavedFile,
  ImageCaptureData,
  AIProviderConfig,
  ToolResponse,
  AIProvider,
} from "../../../src/types";
import * as fs from "fs/promises";
import * as os from "os";
import * as pathModule from "path";

// Mock the Swift CLI utility
vi.mock("../../../src/utils/peekaboo-cli");

// Mock AI Provider utilities
// Declare the variables that will hold the mock functions first
let mockDetermineProviderAndModel: vi.MockedFunction<any>;
let mockAnalyzeImageWithProvider: vi.MockedFunction<any>;
let mockParseAIProviders: vi.MockedFunction<any>;
let mockIsProviderAvailable: vi.MockedFunction<any>;
let mockGetDefaultModelForProvider: vi.MockedFunction<any>;

vi.mock("../../../src/utils/ai-providers", () => {
  // Create new vi.fn() instances inside the factory
  const determineProviderAndModel = vi.fn();
  const analyzeImageWithProvider = vi.fn();
  const parseAIProviders = vi.fn();
  const isProviderAvailable = vi.fn();
  const getDefaultModelForProvider = vi.fn().mockReturnValue("default-model");

  // Assign them to the outer scope variables so tests can reference them
  // This assignment happens AFTER the vi.mock call is processed by Vitest due to hoisting.
  // We will re-assign these correctly after the mock call using an import.
  return {
    determineProviderAndModel,
    analyzeImageWithProvider,
    parseAIProviders,
    isProviderAvailable,
    getDefaultModelForProvider,
  };
});

// Mock fs/promises for mkdtemp, unlink, rmdir
vi.mock("fs/promises");

// Now, import the mocked module and assign the vi.fn() instances to our variables
// This ensures our variables hold the actual mocks created by Vitest's factory.
import * as ActualAiProvidersMock from "../../../src/utils/ai-providers";
mockDetermineProviderAndModel =
  ActualAiProvidersMock.determineProviderAndModel as vi.MockedFunction<any>;
mockAnalyzeImageWithProvider =
  ActualAiProvidersMock.analyzeImageWithProvider as vi.MockedFunction<any>;
mockParseAIProviders =
  ActualAiProvidersMock.parseAIProviders as vi.MockedFunction<any>;
mockIsProviderAvailable =
  ActualAiProvidersMock.isProviderAvailable as vi.MockedFunction<any>;
mockGetDefaultModelForProvider =
  ActualAiProvidersMock.getDefaultModelForProvider as vi.MockedFunction<any>;

const mockExecuteSwiftCli = executeSwiftCli as vi.MockedFunction<
  typeof executeSwiftCli
>;
const mockReadImageAsBase64 = readImageAsBase64 as vi.MockedFunction<
  typeof readImageAsBase64
>;

const mockFsMkdtemp = fs.mkdtemp as vi.MockedFunction<typeof fs.mkdtemp>;
const mockFsUnlink = fs.unlink as vi.MockedFunction<typeof fs.unlink>;
const mockFsRmdir = fs.rmdir as vi.MockedFunction<typeof fs.rmdir>;

const mockLogger = pino({ level: "silent" });
const mockContext = { logger: mockLogger };

const MOCK_TEMP_DIR_PATH = "/tmp/peekaboo-img-mock";
const MOCK_TEMP_IMAGE_PATH = `${MOCK_TEMP_DIR_PATH}/capture.png`;

describe("Image Tool", () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mockFsMkdtemp.mockResolvedValue(MOCK_TEMP_DIR_PATH);
    mockFsUnlink.mockResolvedValue(undefined);
    mockFsRmdir.mockResolvedValue(undefined);
    process.env.PEEKABOO_AI_PROVIDERS = "";

    // Ensure specific mock implementations are reset/re-set for each test or suite as needed
    // The functions themselves are already vi.fn() instances.
    mockDetermineProviderAndModel.mockReset();
    mockAnalyzeImageWithProvider.mockReset();
    mockParseAIProviders.mockReset();
    mockIsProviderAvailable.mockReset();
    mockGetDefaultModelForProvider.mockReset().mockReturnValue("default-model"); // Re-apply default mock behavior if any
  });

  describe("imageToolHandler - Capture Only", () => {
    it("should capture screen with minimal parameters", async () => {
      const mockResponse = mockSwiftCli.captureImage("screen", {});
      mockExecuteSwiftCli.mockResolvedValue(mockResponse);

      const result = await imageToolHandler(
        {
          format: "png",
        return_data: false,
          capture_focus: "background",
        },
        mockContext,
      );

      expect(result.content[0].type).toBe("text");
      expect(result.content[0].text).toContain("Captured 1 image");
      expect(mockExecuteSwiftCli).toHaveBeenCalledWith(
        expect.arrayContaining(["image", "--mode", "screen"]),
        mockLogger,
      );
      expect(result.saved_files).toEqual(mockResponse.data?.saved_files);
      expect(result.analysis_text).toBeUndefined();
      expect(result.model_used).toBeUndefined();
    });

    it("should return image data when return_data is true and no question is asked", async () => {
      const mockSavedFile: SavedFile = {
        path: "/tmp/test.png",
        mime_type: "image/png",
        item_label: "Screen 1",
      };
      const mockCaptureData: ImageCaptureData = {
        saved_files: [mockSavedFile],
      };
      const mockCliResponse = {
        success: true,
        data: mockCaptureData,
        messages: ["Captured one file"],
      };
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);
      mockReadImageAsBase64.mockResolvedValue("base64imagedata");

      const result = await imageToolHandler(
        {
          format: "png",
          return_data: true,
          capture_focus: "background",
        },
        mockContext,
      );

      expect(result.isError).toBeUndefined();
      expect(result.content).toEqual(
        expect.arrayContaining([
          expect.objectContaining({
            type: "text",
            text: expect.stringContaining("Captured 1 image"),
          }),
          expect.objectContaining({ type: "image", data: "base64imagedata" }),
        ]),
      );
      expect(mockReadImageAsBase64).toHaveBeenCalledWith("/tmp/test.png");
      expect(result.saved_files).toEqual([mockSavedFile]);
      expect(result.analysis_text).toBeUndefined();
    });
  });

  describe("imageToolHandler - Capture and Analyze", () => {
    const MOCK_QUESTION = "What is in this image?";
    const MOCK_ANALYSIS_RESPONSE = "This is a cat.";
    const MOCK_PROVIDER_DETAILS: AIProvider = {
      provider: "ollama",
      model: "llava:custom",
    };

    beforeEach(() => {
      mockParseAIProviders.mockReturnValue([
        { provider: "ollama", model: "llava:default" },
      ]);
      mockDetermineProviderAndModel.mockResolvedValue(MOCK_PROVIDER_DETAILS);
      mockAnalyzeImageWithProvider.mockResolvedValue(MOCK_ANALYSIS_RESPONSE);
      mockReadImageAsBase64.mockResolvedValue("base64dataforanalysis");
      process.env.PEEKABOO_AI_PROVIDERS = "ollama/llava:default";
    });

    it("should capture, analyze, and delete temp image if no path provided", async () => {
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: MOCK_TEMP_IMAGE_PATH,
        format: "png",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      const result = await imageToolHandler(
        {
          question: MOCK_QUESTION,
          format: "png",
        },
        mockContext,
      );

      expect(mockFsMkdtemp).toHaveBeenCalled();
      expect(mockExecuteSwiftCli).toHaveBeenCalledWith(
        expect.arrayContaining(["--path", MOCK_TEMP_IMAGE_PATH]),
        mockLogger,
      );
      expect(mockReadImageAsBase64).toHaveBeenCalledWith(MOCK_TEMP_IMAGE_PATH);
      expect(mockDetermineProviderAndModel).toHaveBeenCalled();
      expect(mockAnalyzeImageWithProvider).toHaveBeenCalledWith(
        MOCK_PROVIDER_DETAILS,
        MOCK_TEMP_IMAGE_PATH,
        "base64dataforanalysis",
        MOCK_QUESTION,
        mockLogger,
      );

      expect(result.analysis_text).toBe(MOCK_ANALYSIS_RESPONSE);
      expect(result.model_used).toBe(
        `${MOCK_PROVIDER_DETAILS.provider}/${MOCK_PROVIDER_DETAILS.model}`,
      );
      expect(result.content).toEqual(
        expect.arrayContaining([
          expect.objectContaining({
            text: expect.stringContaining("Captured 1 image"),
          }),
          expect.objectContaining({
            text: expect.stringContaining("Analysis succeeded"),
          }),
          expect.objectContaining({
            text: `Analysis Result: ${MOCK_ANALYSIS_RESPONSE}`,
          }),
        ]),
      );
      expect(result.saved_files).toEqual([]);
      expect(
        result.content.some((item) => item.type === "image" && item.data),
      ).toBe(false);
      expect(mockFsUnlink).toHaveBeenCalledWith(MOCK_TEMP_IMAGE_PATH);
      expect(mockFsRmdir).toHaveBeenCalledWith(MOCK_TEMP_DIR_PATH);
      expect(result.isError).toBeUndefined();
    });

    it("should capture, analyze, and keep image if path IS provided", async () => {
      const USER_PATH = "/user/specified/path.jpg";
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: USER_PATH,
        format: "jpg",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      const result = await imageToolHandler(
        {
          path: USER_PATH,
          question: MOCK_QUESTION,
          format: "jpg",
        },
        mockContext,
      );

      expect(mockFsMkdtemp).not.toHaveBeenCalled();
      expect(mockExecuteSwiftCli).toHaveBeenCalledWith(
        expect.arrayContaining(["--path", USER_PATH]),
        mockLogger,
      );
      expect(mockReadImageAsBase64).toHaveBeenCalledWith(USER_PATH);
      expect(mockAnalyzeImageWithProvider).toHaveBeenCalled();

      expect(result.analysis_text).toBe(MOCK_ANALYSIS_RESPONSE);
      expect(result.saved_files).toEqual(mockCliResponse.data?.saved_files);
      expect(mockFsUnlink).not.toHaveBeenCalled();
      expect(mockFsRmdir).not.toHaveBeenCalled();
      expect(result.isError).toBeUndefined();
    });

    it("should use provider_config if specified", async () => {
      const specificProviderConfig: AIProviderConfig = {
        type: "openai",
        model: "gpt-4-vision",
      };
      const specificProviderDetails: AIProvider = {
        provider: "openai",
        model: "gpt-4-vision",
      };
      mockDetermineProviderAndModel.mockResolvedValue(specificProviderDetails);
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: MOCK_TEMP_IMAGE_PATH,
        format: "png",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      await imageToolHandler(
        {
          question: MOCK_QUESTION,
          provider_config: specificProviderConfig,
          format: "png",
        },
        mockContext,
      );

      expect(mockDetermineProviderAndModel).toHaveBeenCalledWith(
        specificProviderConfig,
        expect.any(Array),
        mockLogger,
      );
      expect(mockAnalyzeImageWithProvider).toHaveBeenCalledWith(
        specificProviderDetails,
        MOCK_TEMP_IMAGE_PATH,
        "base64dataforanalysis",
        MOCK_QUESTION,
        mockLogger,
      );
    });

    it("should handle failure in readImageAsBase64 before analysis", async () => {
      mockReadImageAsBase64.mockRejectedValue(
        new Error("Failed to read image"),
      );
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: MOCK_TEMP_IMAGE_PATH,
        format: "png",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      const result = await imageToolHandler(
        { question: MOCK_QUESTION, format: "png" },
        mockContext,
      );

      expect(mockAnalyzeImageWithProvider).not.toHaveBeenCalled();
      expect(result.analysis_text).toContain(
        "Analysis skipped: Failed to read captured image",
      );
      expect(result.isError).toBe(true);
      expect(result.model_used).toBeUndefined();
      expect(mockFsUnlink).toHaveBeenCalledWith(MOCK_TEMP_IMAGE_PATH);
    });

    it("should handle failure in determineProviderAndModel (rejected promise)", async () => {
      mockDetermineProviderAndModel.mockRejectedValue(
        new Error("No provider available error from determine"),
      );
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: MOCK_TEMP_IMAGE_PATH,
        format: "png",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      const result = await imageToolHandler(
        { question: MOCK_QUESTION, format: "png" },
        mockContext,
      );

      expect(mockAnalyzeImageWithProvider).not.toHaveBeenCalled();
      expect(result.analysis_text).toContain(
        "AI analysis failed: No provider available error from determine",
      );
      expect(result.isError).toBe(true);
    });

    it("should handle failure when determineProviderAndModel resolves to no provider", async () => {
      mockDetermineProviderAndModel.mockResolvedValue({
        provider: null,
        model: "",
      });
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: MOCK_TEMP_IMAGE_PATH,
        format: "png",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      const result = await imageToolHandler(
        { question: MOCK_QUESTION, format: "png" },
        mockContext,
      );

      expect(mockAnalyzeImageWithProvider).not.toHaveBeenCalled();
      expect(result.analysis_text).toContain(
        "Analysis skipped: No AI providers are currently operational",
      );
      expect(result.isError).toBe(true);
    });

    it("should handle failure in analyzeImageWithProvider", async () => {
      mockAnalyzeImageWithProvider.mockRejectedValue(
        new Error("AI API Error from analyze"),
      );
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: MOCK_TEMP_IMAGE_PATH,
        format: "png",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      const result = await imageToolHandler(
        { question: MOCK_QUESTION, format: "png" },
        mockContext,
      );

      expect(result.analysis_text).toContain(
        "AI analysis failed: AI API Error from analyze",
      );
      expect(result.isError).toBe(true);
      expect(result.model_used).toBeUndefined();
    });

    it("should correctly report error if PEEKABOO_AI_PROVIDERS is not set and no provider_config given", async () => {
      process.env.PEEKABOO_AI_PROVIDERS = "";
      mockParseAIProviders.mockReturnValue([]);
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: MOCK_TEMP_IMAGE_PATH,
        format: "png",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      const result = await imageToolHandler(
        { question: MOCK_QUESTION, format: "png" },
        mockContext,
      );

      expect(result.analysis_text).toContain(
        "Analysis skipped: AI analysis not configured on this server",
      );
      expect(result.isError).toBe(true);
      expect(mockAnalyzeImageWithProvider).not.toHaveBeenCalled();
    });

    it("should return isError = true if analysis is attempted but fails, even if capture succeeds", async () => {
      mockAnalyzeImageWithProvider.mockRejectedValue(new Error("AI Error"));
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: MOCK_TEMP_IMAGE_PATH,
        format: "png",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      const result = await imageToolHandler(
        { question: MOCK_QUESTION, format: "png" },
        mockContext,
      );
      expect(result.isError).toBe(true);
      expect(result.content[0].text).toContain("Captured 1 image");
      expect(result.content[0].text).toContain("Analysis failed/skipped");
      expect(result.analysis_text).toContain("AI analysis failed: AI Error");
    });

    it("should NOT return base64_data in content if question is asked, even if return_data is true", async () => {
      const mockCliResponse = mockSwiftCli.captureImage("screen", {
        path: MOCK_TEMP_IMAGE_PATH,
        format: "png",
      });
      mockExecuteSwiftCli.mockResolvedValue(mockCliResponse);

      const result = await imageToolHandler(
        {
          question: MOCK_QUESTION,
        return_data: true,
          format: "png",
        },
        mockContext,
      );

      expect(
        result.content.some((item) => item.type === "image" && item.data),
      ).toBe(false);
      expect(result.analysis_text).toBe(MOCK_ANALYSIS_RESPONSE);
    });
  });

  describe("buildSwiftCliArgs", () => {
    const defaults = {
      format: "png" as const,
      return_data: false,
      capture_focus: "background" as const,
    };

    it("should default to screen mode if no app provided and no mode specified", () => {
      const args = buildSwiftCliArgs({ ...defaults });
      expect(args).toEqual([
        "image",
        "--mode",
        "screen",
        "--format",
        "png",
        "--capture-focus",
        "background",
      ]);
    });

    it("should default to window mode if app is provided and no mode specified", () => {
      const args = buildSwiftCliArgs({ ...defaults, app: "Safari" });
      expect(args).toEqual([
        "image",
        "--app",
        "Safari",
        "--mode",
        "window",
        "--format",
        "png",
        "--capture-focus",
        "background",
      ]);
    });

    it("should use specified mode: screen", () => {
      const args = buildSwiftCliArgs({ ...defaults, mode: "screen" });
      expect(args).toEqual(expect.arrayContaining(["--mode", "screen"]));
    });

    it("should use specified mode: window with app", () => {
      const args = buildSwiftCliArgs({
        ...defaults,
        app: "Terminal",
        mode: "window",
      });
      expect(args).toEqual(
        expect.arrayContaining(["--app", "Terminal", "--mode", "window"]),
      );
    });

    it("should use specified mode: multi with app", () => {
      const args = buildSwiftCliArgs({
        ...defaults,
        app: "Finder",
        mode: "multi",
      });
      expect(args).toEqual(
        expect.arrayContaining(["--app", "Finder", "--mode", "multi"]),
      );
    });

    it("should include app", () => {
      const args = buildSwiftCliArgs({ ...defaults, app: "Notes" });
      expect(args).toEqual(expect.arrayContaining(["--app", "Notes"]));
    });

    it("should include path", () => {
      const args = buildSwiftCliArgs({ ...defaults, path: "/tmp/image.jpg" });
      expect(args).toEqual(
        expect.arrayContaining(["--path", "/tmp/image.jpg"]),
      );
    });

    it("should include window_specifier by title", () => {
      const args = buildSwiftCliArgs({
        ...defaults,
        app: "Safari",
        window_specifier: { title: "Apple" },
      });
      expect(args).toEqual(expect.arrayContaining(["--window-title", "Apple"]));
    });

    it("should include window_specifier by index", () => {
      const args = buildSwiftCliArgs({
        ...defaults,
        app: "Safari",
        window_specifier: { index: 0 },
      });
      expect(args).toEqual(expect.arrayContaining(["--window-index", "0"]));
    });

    it("should include format (default png)", () => {
      const args = buildSwiftCliArgs({ ...defaults });
      expect(args).toEqual(expect.arrayContaining(["--format", "png"]));
    });

    it("should include specified format jpg", () => {
      const args = buildSwiftCliArgs({ ...defaults, format: "jpg" });
      expect(args).toEqual(expect.arrayContaining(["--format", "jpg"]));
    });

    it("should include capture_focus (default background)", () => {
      const args = buildSwiftCliArgs({ ...defaults });
      expect(args).toEqual(
        expect.arrayContaining(["--capture-focus", "background"]),
      );
    });

    it("should include specified capture_focus foreground", () => {
      const args = buildSwiftCliArgs({
        ...defaults,
        capture_focus: "foreground",
      });
      expect(args).toEqual(
        expect.arrayContaining(["--capture-focus", "foreground"]),
      );
    });

    it("should handle all options together", () => {
      const input: ImageToolInput = {
        ...defaults, // Ensure all required fields are present
        app: "Preview",
        path: "/users/test/file.tiff",
        mode: "window",
        window_specifier: { index: 1 },
        format: "png",
        capture_focus: "foreground",
      };
      const args = buildSwiftCliArgs(input);
      expect(args).toEqual([
        "image",
        "--app",
        "Preview",
        "--path",
        "/users/test/file.tiff",
        "--mode",
        "window",
        "--window-index",
        "1",
        "--format",
        "png",
        "--capture-focus",
        "foreground",
      ]);
    });

    it("should use input.path if provided, even with a question", () => {
      const input: ImageToolInput = { path: "/my/path.png", question: "test" };
      const args = buildSwiftCliArgs(input);
      expect(args).toContain("--path");
      expect(args).toContain("/my/path.png");
    });

    it("should NOT use PEEKABOO_DEFAULT_SAVE_PATH if a question is asked", () => {
      process.env.PEEKABOO_DEFAULT_SAVE_PATH = "/default/env.png";
      const input: ImageToolInput = { question: "test" };
      const args = buildSwiftCliArgs(input);
      expect(args.includes("--path")).toBe(false);
      delete process.env.PEEKABOO_DEFAULT_SAVE_PATH;
    });

    it("should use PEEKABOO_DEFAULT_SAVE_PATH if no path and no question", () => {
      process.env.PEEKABOO_DEFAULT_SAVE_PATH = "/default/env.png";
      const input: ImageToolInput = {};
      const args = buildSwiftCliArgs(input);
      expect(args).toContain("--path");
      expect(args).toContain("/default/env.png");
      delete process.env.PEEKABOO_DEFAULT_SAVE_PATH;
    });

    it("should use default format and capture_focus if not provided", () => {
      const input: ImageToolInput = {
        format: "png",
        capture_focus: "background",
      };
      const args = buildSwiftCliArgs(input);
      expect(args).toContain("--format");
      expect(args).toContain("png");
      expect(args).toContain("--capture-focus");
      expect(args).toContain("background");
    });
  });
});