conversationdb

import { type PluginContext, tool } from "@lmstudio/sdk";
import { z } from "zod";
import { ConversationDB, Message } from "./dbManager";

const db = new ConversationDB();

// --- Conversation Tools ---

const saveConversation = tool({
  name: "save_conversation",
  description:
    "Save a conversation with messages. Provide a title, array of messages (each with role, content, timestamp), and optional tags and summary.",
  parameters: {
    title: z.string().describe("Title for this conversation"),
    messages: z
      .array(
        z.object({
          role: z.enum(["system", "user", "assistant"]).describe("Message role"),
          content: z.string().describe("Message content"),
          timestamp: z.string().describe("ISO timestamp or human-readable time"),
        })
      )
      .describe("Array of messages in the conversation"),
    tags: z.array(z.string()).optional().describe("Tags to categorize this conversation"),
    summary: z.string().optional().describe("Brief summary of the conversation"),
  },
  implementation: async ({ title, messages, tags, summary }) => {
    return db.addConversation(title, messages as Message[], tags || [], summary || "");
  },
});

const getConversation = tool({
  name: "get_conversation",
  description: "Retrieve a saved conversation by ID. Set includeMessages to true to see full message history.",
  parameters: {
    id: z.number().describe("Conversation ID"),
    includeMessages: z.boolean().optional().default(true).describe("Include full message history"),
  },
  implementation: async ({ id, includeMessages }) => {
    return db.getConversation(id, includeMessages);
  },
});

const listConversations = tool({
  name: "list_conversations",
  description: "List saved conversations, optionally filtered by tag.",
  parameters: {
    tag: z.string().optional().describe("Filter by tag (partial match)"),
    limit: z.number().optional().default(20).describe("Maximum number of results"),
  },
  implementation: async ({ tag, limit }) => {
    return db.listConversations(tag, limit);
  },
});

const searchConversations = tool({
  name: "search_conversations",
  description: "Search conversations by query string (matches title, summary, tags, and message content).",
  parameters: {
    query: z.string().describe("Search query"),
    limit: z.number().optional().default(10).describe("Maximum number of results"),
  },
  implementation: async ({ query, limit }) => {
    return db.searchConversations(query, limit);
  },
});

const updateConversation = tool({
  name: "update_conversation",
  description: "Update a conversation's title, tags, or summary by ID.",
  parameters: {
    id: z.number().describe("Conversation ID"),
    title: z.string().optional().describe("New title"),
    tags: z.array(z.string()).optional().describe("New tags (replaces existing)"),
    summary: z.string().optional().describe("New summary"),
  },
  implementation: async ({ id, title, tags, summary }) => {
    const updates: { title?: string; tags?: string[]; summary?: string } = {};
    if (title !== undefined) updates.title = title;
    if (tags !== undefined) updates.tags = tags;
    if (summary !== undefined) updates.summary = summary;
    return db.updateConversation(id, updates);
  },
});

const deleteConversation = tool({
  name: "delete_conversation",
  description: "Permanently delete a conversation and all its associated images.",
  parameters: {
    id: z.number().describe("Conversation ID to delete"),
  },
  implementation: async ({ id }) => {
    return db.deleteConversation(id);
  },
});

// --- Image Tools ---

const saveImage = tool({
  name: "save_image",
  description:
    "Save an image to a conversation. Provide the conversation ID, filename, description, and optionally base64 image data.",
  parameters: {
    conversationId: z.number().describe("ID of the conversation to attach this image to"),
    filename: z.string().describe("Filename for the image (e.g., screenshot.png)"),
    description: z.string().describe("Description of what the image shows"),
    base64Data: z.string().optional().describe("Base64-encoded image data (with or without data: prefix)"),
  },
  implementation: async ({ conversationId, filename, description, base64Data }) => {
    return db.addImage(conversationId, filename, description, base64Data);
  },
});

const getImage = tool({
  name: "get_image",
  description: "Retrieve image metadata by ID.",
  parameters: {
    id: z.number().describe("Image ID"),
    includeData: z.boolean().optional().default(false).describe("Include base64 data length info"),
  },
  implementation: async ({ id, includeData }) => {
    return db.getImage(id, includeData);
  },
});

const listImages = tool({
  name: "list_images",
  description: "List all images, optionally filtered by conversation ID.",
  parameters: {
    conversationId: z.number().optional().describe("Filter by conversation ID"),
  },
  implementation: async ({ conversationId }) => {
    return db.listImages(conversationId);
  },
});

const deleteImage = tool({
  name: "delete_image",
  description: "Permanently delete an image by ID.",
  parameters: {
    id: z.number().describe("Image ID to delete"),
  },
  implementation: async ({ id }) => {
    return db.deleteImage(id);
  },
});

// --- Utility Tools ---

const dbStats = tool({
  name: "conversation_db_stats",
  description: "Show database statistics: conversation count, message count, image count, and top tags.",
  parameters: {},
  implementation: async () => {
    return db.stats();
  },
});

const clearAll = tool({
  name: "clear_all_conversations",
  description: "DANGER: Delete ALL conversations and images. Use only when explicitly asked to wipe the database.",
  parameters: {},
  implementation: async () => {
    return db.clearAll();
  },
});

// --- Main Entry Point ---

export async function main(context: PluginContext) {
  context.withToolsProvider(async () => {
    return [
      saveConversation,
      getConversation,
      listConversations,
      searchConversations,
      updateConversation,
      deleteConversation,
      saveImage,
      getImage,
      listImages,
      deleteImage,
      dbStats,
      clearAll,
    ];
  });
}