import { tool, type Tool, type ToolsProvider } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics, resolveDatabasePath } from "./config";
import { KnowledgeGraphManager } from "./knowledgeGraphManager";
import { EntitySchema, RelationSchema, SearchNodesQuerySchema } from "./schemas";

export const toolsProvider: ToolsProvider = async ctl => {
  const pluginConfig = ctl.getPluginConfig(configSchematics);
  const databasePath = resolveDatabasePath(pluginConfig.get("databasePath"));
  const manager = new KnowledgeGraphManager(databasePath);

  const tools: Tool[] = [];

  tools.push(
    tool({
      name: "create_entities",
      description: "Create multiple new entities in the knowledge graph",
      parameters: { entities: z.array(EntitySchema) },
      implementation: async ({ entities }) => manager.createEntities(entities),
    }),
  );

  tools.push(
    tool({
      name: "create_relations",
      description:
        "Create multiple new relations between entities in the knowledge graph. Relations should be in active voice",
      parameters: { relations: z.array(RelationSchema) },
      implementation: async ({ relations }) => manager.createRelations(relations),
    }),
  );

  tools.push(
    tool({
      name: "add_observations",
      description: "Add new observations to existing entities in the knowledge graph",
      parameters: {
        observations: z.array(
          z.object({
            entityName: z.string().describe("The name of the entity to add the observations to"),
            contents: z.array(z.string()).describe("An array of observation contents to add"),
          }),
        ),
      },
      implementation: async ({ observations }) => manager.addObservations(observations),
    }),
  );

  tools.push(
    tool({
      name: "delete_entities",
      description: "Delete multiple entities and their associated relations from the knowledge graph",
      parameters: { entityNames: z.array(z.string()).describe("An array of entity names to delete") },
      implementation: async ({ entityNames }) => {
        const { deleted, notFound } = await manager.deleteEntities(entityNames);
        return notFound.length === 0
          ? "Entities deleted successfully"
          : `Deleted ${deleted.length} of ${entityNames.length} entities. Not found: ${notFound.join(", ")}`;
      },
    }),
  );

  tools.push(
    tool({
      name: "delete_observations",
      description: "Delete specific observations from entities in the knowledge graph",
      parameters: {
        deletions: z.array(
          z.object({
            entityName: z.string().describe("The name of the entity containing the observations"),
            observations: z.array(z.string()).describe("An array of observations to delete"),
          }),
        ),
      },
      implementation: async ({ deletions }) => {
        const { deletedCount, missingEntities } = await manager.deleteObservations(deletions);
        const requested = deletions.reduce((total, d) => total + d.observations.length, 0);
        return deletedCount === requested
          ? "Observations deleted successfully"
          : `Deleted ${deletedCount} of ${requested} observations.` +
              (missingEntities.length ? ` Entities not found: ${missingEntities.join(", ")}` : "");
      },
    }),
  );

  tools.push(
    tool({
      name: "delete_relations",
      description: "Delete multiple relations from the knowledge graph",
      parameters: { relations: z.array(RelationSchema).describe("An array of relations to delete") },
      implementation: async ({ relations }) => {
        const { deletedCount } = await manager.deleteRelations(relations);
        return deletedCount === relations.length
          ? "Relations deleted successfully"
          : `Deleted ${deletedCount} of ${relations.length} relations. The rest matched nothing.`;
      },
    }),
  );

  tools.push(
    tool({
      name: "read_graph",
      description: "Read the entire knowledge graph",
      parameters: {},
      implementation: async () => manager.readGraph(),
    }),
  );

  tools.push(
    tool({
      name: "search_nodes",
      description: "Search for nodes in the knowledge graph based on a query",
      parameters: { query: SearchNodesQuerySchema },
      implementation: async ({ query }) => manager.searchNodes(query),
    }),
  );

  tools.push(
    tool({
      name: "open_nodes",
      description: "Open specific nodes in the knowledge graph by their names",
      parameters: { names: z.array(z.string()).describe("An array of entity names to retrieve") },
      implementation: async ({ names }) => manager.openNodes(names),
    }),
  );

  return tools;
};

import { tool, type Tool, type ToolsProvider } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics, resolveDatabasePath } from "./config";
import { KnowledgeGraphManager } from "./knowledgeGraphManager";
import { EntitySchema, RelationSchema, SearchNodesQuerySchema } from "./schemas";

export const toolsProvider: ToolsProvider = async ctl => {
  const pluginConfig = ctl.getPluginConfig(configSchematics);
  const databasePath = resolveDatabasePath(pluginConfig.get("databasePath"));
  const manager = new KnowledgeGraphManager(databasePath);

  const tools: Tool[] = [];

  tools.push(
    tool({
      name: "create_entities",
      description: "Create multiple new entities in the knowledge graph",
      parameters: { entities: z.array(EntitySchema) },
      implementation: async ({ entities }) => manager.createEntities(entities),
    }),
  );

  tools.push(
    tool({
      name: "create_relations",
      description:
        "Create multiple new relations between entities in the knowledge graph. Relations should be in active voice",
      parameters: { relations: z.array(RelationSchema) },
      implementation: async ({ relations }) => manager.createRelations(relations),
    }),
  );

  tools.push(
    tool({
      name: "add_observations",
      description: "Add new observations to existing entities in the knowledge graph",
      parameters: {
        observations: z.array(
          z.object({
            entityName: z.string().describe("The name of the entity to add the observations to"),
            contents: z.array(z.string()).describe("An array of observation contents to add"),
          }),
        ),
      },
      implementation: async ({ observations }) => manager.addObservations(observations),
    }),
  );

  tools.push(
    tool({
      name: "delete_entities",
      description: "Delete multiple entities and their associated relations from the knowledge graph",
      parameters: { entityNames: z.array(z.string()).describe("An array of entity names to delete") },
      implementation: async ({ entityNames }) => {
        const { deleted, notFound } = await manager.deleteEntities(entityNames);
        return notFound.length === 0
          ? "Entities deleted successfully"
          : `Deleted ${deleted.length} of ${entityNames.length} entities. Not found: ${notFound.join(", ")}`;
      },
    }),
  );

  tools.push(
    tool({
      name: "delete_observations",
      description: "Delete specific observations from entities in the knowledge graph",
      parameters: {
        deletions: z.array(
          z.object({
            entityName: z.string().describe("The name of the entity containing the observations"),
            observations: z.array(z.string()).describe("An array of observations to delete"),
          }),
        ),
      },
      implementation: async ({ deletions }) => {
        const { deletedCount, missingEntities } = await manager.deleteObservations(deletions);
        const requested = deletions.reduce((total, d) => total + d.observations.length, 0);
        return deletedCount === requested
          ? "Observations deleted successfully"
          : `Deleted ${deletedCount} of ${requested} observations.` +
              (missingEntities.length ? ` Entities not found: ${missingEntities.join(", ")}` : "");
      },
    }),
  );

  tools.push(
    tool({
      name: "delete_relations",
      description: "Delete multiple relations from the knowledge graph",
      parameters: { relations: z.array(RelationSchema).describe("An array of relations to delete") },
      implementation: async ({ relations }) => {
        const { deletedCount } = await manager.deleteRelations(relations);
        return deletedCount === relations.length
          ? "Relations deleted successfully"
          : `Deleted ${deletedCount} of ${relations.length} relations. The rest matched nothing.`;
      },
    }),
  );

  tools.push(
    tool({
      name: "read_graph",
      description: "Read the entire knowledge graph",
      parameters: {},
      implementation: async () => manager.readGraph(),
    }),
  );

  tools.push(
    tool({
      name: "search_nodes",
      description: "Search for nodes in the knowledge graph based on a query",
      parameters: { query: SearchNodesQuerySchema },
      implementation: async ({ query }) => manager.searchNodes(query),
    }),
  );

  tools.push(
    tool({
      name: "open_nodes",
      description: "Open specific nodes in the knowledge graph by their names",
      parameters: { names: z.array(z.string()).describe("An array of entity names to retrieve") },
      implementation: async ({ names }) => manager.openNodes(names),
    }),
  );

  return tools;
};