data-plugin

import { tool } from "@lmstudio/sdk";
import { z } from "zod";
import { v4 as uuidv4 } from "uuid";
import * as csvEngine from "./engine/csv-engine";
import * as schemaManager from "./engine/schema";
import * as indexManager from "./engine/index-manager";
import { getDataRoot, initDataRoot } from "./config";
import {
  DataPluginError,
  NotFoundError,
  DuplicateError,
  ValidationError,
} from "./errors";
import type { ColumnDef, TableSchema } from "./types";

const ColumnDefSchema = z.object({
  name: z.string().min(1),
  type: z.enum(["string", "number", "boolean", "date"]),
  required: z.boolean().optional(),
  primaryKey: z.boolean().optional(),
  unique: z.boolean().optional(),
  defaultValue: z.unknown().optional(),
  description: z.string().optional(),
});

async function ensureDataRoot(): Promise<string> {
  const root = getDataRoot();
  await initDataRoot(root);
  return root;
}

export async function createToolsProvider() {
  const dataRoot = await ensureDataRoot();

  return [
    // ── create_database ──────────────────────────────────────────────
    tool({
      name: "create_database",
      description:
        "Create a new database (a directory that will contain tables)",
      parameters: {
        name: z.string().min(1).describe("Database name"),
        displayName: z.string().optional().describe("Optional display name"),
        description: z.string().optional().describe("Optional description"),
      },
      implementation: async ({ name, displayName, description }, ctx) => {
        try {
          await indexManager.initDatabase(
            dataRoot,
            name,
            displayName,
            description,
          );
          return JSON.stringify({
            success: true,
            name,
            path: `${dataRoot}/databases/${name}`,
          });
        } catch (e) {
          if (e instanceof DuplicateError)
            return JSON.stringify({ success: false, error: e.message });
          throw e;
        }
      },
    }),

    // ── delete_database ──────────────────────────────────────────────
    tool({
      name: "delete_database",
      description: "Delete a database and all its tables permanently",
      parameters: {
        name: z.string().min(1).describe("Database name to delete"),
        confirm: z.boolean().describe("Must be true to confirm deletion"),
      },
      implementation: async ({ name, confirm }, ctx) => {
        if (!confirm)
          return JSON.stringify({
            success: false,
            error: "Deletion not confirmed",
          });
        try {
          const idx = await indexManager.getIndex(dataRoot, name);
          const tableCount = Object.keys(idx.tables).length;
          await indexManager.deleteDatabase(dataRoot, name);
          return JSON.stringify({ success: true, deletedTables: tableCount });
        } catch (e) {
          if (e instanceof NotFoundError)
            return JSON.stringify({ success: false, error: e.message });
          throw e;
        }
      },
    }),

    // ── list_databases ───────────────────────────────────────────────
    tool({
      name: "list_databases",
      description: "List all databases with summary information",
      parameters: {},
      implementation: async () => {
        const dbs = await indexManager.getDatabases(dataRoot);
        return JSON.stringify({ databases: dbs, total: dbs.length });
      },
    }),

    // ── get_database_info ────────────────────────────────────────────
    tool({
      name: "get_database_info",
      description:
        "Get detailed information about a database including its tables",
      parameters: {
        name: z.string().min(1).describe("Database name"),
      },
      implementation: async ({ name }) => {
        const idx = await indexManager.getIndex(dataRoot, name);
        const tables = await indexManager.getTables(dataRoot, name);
        const tableCount = Object.keys(idx.tables).length;
        const totalSizeBytes = tables.reduce((s, t) => s + t.sizeBytes, 0);
        return JSON.stringify({
          name,
          displayName: idx.displayName,
          tableCount,
          totalSizeBytes,
          createdAt: idx.createdAt,
          updatedAt: idx.updatedAt,
          tables,
        });
      },
    }),

    // ── create_table ─────────────────────────────────────────────────
    tool({
      name: "create_table",
      description: "Create a new table in a database with a defined schema",
      parameters: {
        database: z.string().min(1).describe("Database name"),
        table: z.string().min(1).describe("Table name"),
        columns: z.array(ColumnDefSchema).min(1).describe("Column definitions"),
        description: z.string().optional(),
      },
      implementation: async ({ database, table, columns, description }) => {
        const exists = await indexManager.tableExists(
          dataRoot,
          database,
          table,
        );
        if (exists) {
          return JSON.stringify({
            success: false,
            error: `Table "${table}" already exists in "${database}"`,
          });
        }

        const schema = schemaManager.createSchema(columns, description);
        const dbDir = `${dataRoot}/databases/${database}`;
        const schemaFile = schemaManager.schemaPath(database, table, dataRoot);

        await schemaManager.saveSchema(schemaFile, schema);
        await csvEngine.writeCSV(`${dbDir}/${table}.csv`, [], columns);

        const idx = await indexManager.getIndex(dataRoot, database);
        idx.tables[table] = {
          name: table,
          displayName: table,
          createdAt: schema.createdAt,
          rowCount: 0,
          sizeBytes: 0,
        };
        await indexManager.updateIndex(dataRoot, database, {
          tables: idx.tables,
        });

        return JSON.stringify({
          success: true,
          database,
          table,
          columnCount: columns.length,
        });
      },
    }),

    // ── delete_table ─────────────────────────────────────────────────
    tool({
      name: "delete_table",
      description: "Delete a table and its schema permanently",
      parameters: {
        database: z.string().min(1).describe("Database name"),
        table: z.string().min(1).describe("Table name"),
        confirm: z.boolean().describe("Must be true to confirm deletion"),
      },
      implementation: async ({ database, table, confirm }) => {
        if (!confirm)
          return JSON.stringify({
            success: false,
            error: "Deletion not confirmed",
          });
        const dbDir = `${dataRoot}/databases/${database}`;

        try {
          const idx = await indexManager.getIndex(dataRoot, database);
          if (!(table in idx.tables)) {
            return JSON.stringify({
              success: false,
              error: `Table "${table}" not found`,
            });
          }
          delete idx.tables[table];
          await indexManager.updateIndex(dataRoot, database, {
            tables: idx.tables,
          });

          await fsRm(`${dbDir}/${table}.csv`);
          await fsRm(`${dbDir}/${table}.schema.json`);

          return JSON.stringify({
            success: true,
            deletedFiles: [`${table}.csv`, `${table}.schema.json`],
          });
        } catch (e) {
          if (e instanceof NotFoundError)
            return JSON.stringify({ success: false, error: e.message });
          throw e;
        }
      },
    }),

    // ── get_schema ───────────────────────────────────────────────────
    tool({
      name: "get_schema",
      description: "Get the schema (column definitions) of a table",
      parameters: {
        database: z.string().min(1),
        table: z.string().min(1),
      },
      implementation: async ({ database, table }) => {
        const schemaFile = schemaManager.schemaPath(database, table, dataRoot);
        const schema = await schemaManager.loadSchema(schemaFile);
        return JSON.stringify({
          database,
          table,
          columns: schema.columns,
          createdAt: schema.createdAt,
          description: schema.description,
        });
      },
    }),

    // ── list_tables ──────────────────────────────────────────────────
    tool({
      name: "list_tables",
      description: "List all tables in a database",
      parameters: {
        database: z.string().min(1).describe("Database name"),
      },
      implementation: async ({ database }) => {
        const tables = await indexManager.getTables(dataRoot, database);
        return JSON.stringify({ database, tables, total: tables.length });
      },
    }),

    // ── insert_row ───────────────────────────────────────────────────
    tool({
      name: "insert_row",
      description:
        "Insert a single row into a table. Returns the generated row ID.",
      parameters: {
        database: z.string().min(1),
        table: z.string().min(1),
        data: z.record(z.unknown()).describe("Row data as key-value pairs"),
      },
      implementation: async ({ database, table, data }) => {
        const schemaFile = schemaManager.schemaPath(database, table, dataRoot);
        const schema = await schemaManager.loadSchema(schemaFile);

        const id = uuidv4();
        const row = { id, ...data };

        const validation = schemaManager.validateRowAgainstSchema(row, schema);
        if (!validation.valid) {
          return JSON.stringify({
            success: false,
            error: `Validation failed: ${validation.errors.map((e) => e.message).join("; ")}`,
          });
        }

        const columns = schema.columns;
        const dbDir = `${dataRoot}/databases/${database}`;
        await csvEngine.appendToCSV(`${dbDir}/${table}.csv`, row, columns);

        const idx = await indexManager.getIndex(dataRoot, database);
        if (idx.tables[table]) {
          idx.tables[table].rowCount = (idx.tables[table].rowCount || 0) + 1;
          await indexManager.updateIndex(dataRoot, database, {
            tables: idx.tables,
          });
        }

        return JSON.stringify({ success: true, id, table });
      },
    }),

    // ── insert_batch ─────────────────────────────────────────────────
    tool({
      name: "insert_batch",
      description: "Insert multiple rows into a table in one operation",
      parameters: {
        database: z.string().min(1),
        table: z.string().min(1),
        rows: z
          .array(z.record(z.unknown()))
          .min(1)
          .describe("Array of row objects"),
      },
      implementation: async ({ database, table, rows }) => {
        const schemaFile = schemaManager.schemaPath(database, table, dataRoot);
        const schema = await schemaManager.loadSchema(schemaFile);

        const enriched = rows.map((r) => ({ id: uuidv4(), ...r }));
        const validation = schemaManager.validateRowsAgainstSchema(
          enriched,
          schema,
        );
        if (!validation.valid) {
          return JSON.stringify({
            success: false,
            error: `Validation failed: ${validation.errors.map((e) => e.message).join("; ")}`,
          });
        }

        const columns = schema.columns;
        const dbDir = `${dataRoot}/databases/${database}`;
        await csvEngine.appendRowsToCSV(
          `${dbDir}/${table}.csv`,
          enriched,
          columns,
        );

        const ids = enriched.map((r) => r.id as string);
        const idx = await indexManager.getIndex(dataRoot, database);
        if (idx.tables[table]) {
          idx.tables[table].rowCount =
            (idx.tables[table].rowCount || 0) + enriched.length;
          await indexManager.updateIndex(dataRoot, database, {
            tables: idx.tables,
          });
        }

        return JSON.stringify({
          success: true,
          inserted: enriched.length,
          ids,
        });
      },
    }),

    // ── select_rows ──────────────────────────────────────────────────
    tool({
      name: "select_rows",
      description:
        "Select rows from a table with optional filtering, limit and offset",
      parameters: {
        database: z.string().min(1),
        table: z.string().min(1),
        filter: z
          .record(z.unknown())
          .optional()
          .describe("Key-value pairs to filter by (exact match)"),
        limit: z
          .number()
          .int()
          .positive()
          .optional()
          .describe("Max rows to return"),
        offset: z
          .number()
          .int()
          .min(0)
          .optional()
          .describe("Number of rows to skip"),
      },
      implementation: async ({ database, table, filter, limit, offset }) => {
        const dbDir = `${dataRoot}/databases/${database}`;
        const allRows = await csvEngine.parseCSV(`${dbDir}/${table}.csv`);

        let filtered = allRows;
        if (filter) {
          for (const [key, value] of Object.entries(filter)) {
            filtered = filtered.filter((r) => String(r[key]) === String(value));
          }
        }

        const total = filtered.length;
        const offsetVal = offset ?? 0;
        const limitVal = limit ?? total;
        const page = filtered.slice(offsetVal, offsetVal + limitVal);

        return JSON.stringify({
          rows: page,
          total: allRows.length,
          filteredTotal: total,
        });
      },
    }),

    // ── update_row ───────────────────────────────────────────────────
    tool({
      name: "update_row",
      description:
        "Update a row by its ID. Only the provided fields will be updated.",
      parameters: {
        database: z.string().min(1),
        table: z.string().min(1),
        id: z.string().min(1).describe("Row UUID to update"),
        data: z.record(z.unknown()).describe("Fields to update"),
      },
      implementation: async ({ database, table, id, data }) => {
        const dbDir = `${dataRoot}/databases/${database}`;
        const csvPath = `${dbDir}/${table}.csv`;
        const schemaFile = schemaManager.schemaPath(database, table, dataRoot);
        const schema = await schemaManager.loadSchema(schemaFile);

        const allRows = await csvEngine.parseCSV(csvPath);
        const idx = allRows.findIndex((r) => String(r.id) === id);
        if (idx === -1) {
          return JSON.stringify({
            success: false,
            error: `Row with id "${id}" not found`,
          });
        }

        const updated = { ...allRows[idx], ...data };
        allRows[idx] = updated;
        await csvEngine.writeCSV(csvPath, allRows, schema.columns);

        return JSON.stringify({
          success: true,
          id,
          updated: Object.keys(data),
        });
      },
    }),

    // ── delete_row ───────────────────────────────────────────────────
    tool({
      name: "delete_row",
      description: "Delete a row by its ID",
      parameters: {
        database: z.string().min(1),
        table: z.string().min(1),
        id: z.string().min(1).describe("Row UUID to delete"),
      },
      implementation: async ({ database, table, id }) => {
        const dbDir = `${dataRoot}/databases/${database}`;
        const csvPath = `${dbDir}/${table}.csv`;
        const schemaFile = schemaManager.schemaPath(database, table, dataRoot);
        const schema = await schemaManager.loadSchema(schemaFile);

        const allRows = await csvEngine.parseCSV(csvPath);
        const filtered = allRows.filter((r) => String(r.id) !== id);
        if (filtered.length === allRows.length) {
          return JSON.stringify({
            success: false,
            error: `Row with id "${id}" not found`,
          });
        }

        await csvEngine.writeCSV(csvPath, filtered, schema.columns);

        const idx = await indexManager.getIndex(dataRoot, database);
        if (idx.tables[table]) {
          idx.tables[table].rowCount = Math.max(
            0,
            (idx.tables[table].rowCount || 1) - 1,
          );
          await indexManager.updateIndex(dataRoot, database, {
            tables: idx.tables,
          });
        }

        return JSON.stringify({ success: true, id });
      },
    }),

    // ── get_row ──────────────────────────────────────────────────────
    tool({
      name: "get_row",
      description: "Get a single row by its ID",
      parameters: {
        database: z.string().min(1),
        table: z.string().min(1),
        id: z.string().min(1).describe("Row UUID"),
      },
      implementation: async ({ database, table, id }) => {
        const dbDir = `${dataRoot}/databases/${database}`;
        const allRows = await csvEngine.parseCSV(`${dbDir}/${table}.csv`);
        const row = allRows.find((r) => String(r.id) === id) ?? null;
        return JSON.stringify({ row });
      },
    }),

    // ── query_table ──────────────────────────────────────────────────
    tool({
      name: "query_table",
      description: "Advanced query with filters, search, sort, and pagination",
      parameters: {
        database: z.string().min(1),
        table: z.string().min(1),
        filters: z
          .array(
            z.object({
              column: z.string(),
              operator: z.enum([
                "eq",
                "neq",
                "gt",
                "gte",
                "lt",
                "lte",
                "contains",
              ]),
              value: z.unknown(),
            }),
          )
          .optional()
          .describe("Array of filter conditions"),
        search: z
          .object({
            column: z.string(),
            query: z.string(),
            caseSensitive: z.boolean().optional(),
          })
          .optional()
          .describe("Full-text search on a column"),
        sort: z
          .object({
            column: z.string(),
            direction: z.enum(["asc", "desc"]),
          })
          .optional()
          .describe("Sort configuration"),
        limit: z.number().int().positive().optional(),
        offset: z.number().int().min(0).optional(),
      },
      implementation: async ({
        database,
        table,
        filters,
        search,
        sort,
        limit,
        offset,
      }) => {
        const dbDir = `${dataRoot}/databases/${database}`;
        const allRows = await csvEngine.parseCSV(`${dbDir}/${table}.csv`);

        let result = [...allRows];

        if (filters) {
          for (const f of filters) {
            result = result.filter((r) => {
              const cell = r[f.column];
              switch (f.operator) {
                case "eq":
                  return String(cell) === String(f.value);
                case "neq":
                  return String(cell) !== String(f.value);
                case "gt":
                  return Number(cell) > Number(f.value);
                case "gte":
                  return Number(cell) >= Number(f.value);
                case "lt":
                  return Number(cell) < Number(f.value);
                case "lte":
                  return Number(cell) <= Number(f.value);
                case "contains":
                  return String(cell)
                    .toLowerCase()
                    .includes(String(f.value).toLowerCase());
                default:
                  return true;
              }
            });
          }
        }

        if (search) {
          result = result.filter((r) => {
            const cell = String(r[search.column] ?? "");
            const q = search.caseSensitive
              ? search.query
              : search.query.toLowerCase();
            const val = search.caseSensitive ? cell : cell.toLowerCase();
            return val.includes(q);
          });
        }

        if (sort) {
          result.sort((a, b) => {
            const aVal = a[sort.column];
            const bVal = b[sort.column];
            if (aVal == null) return 1;
            if (bVal == null) return -1;
            const cmp =
              typeof aVal === "number"
                ? aVal - Number(bVal)
                : String(aVal).localeCompare(String(bVal));
            return sort.direction === "desc" ? -cmp : cmp;
          });
        }

        const total = result.length;
        const offsetVal = offset ?? 0;
        const limitVal = limit ?? total;
        const page = result.slice(offsetVal, offsetVal + limitVal);
        const totalPages = limitVal > 0 ? Math.ceil(total / limitVal) : 1;

        return JSON.stringify({
          rows: page,
          total: allRows.length,
          filteredTotal: total,
          page: Math.floor(offsetVal / (limitVal || 1)) + 1,
          totalPages,
        });
      },
    }),

    // ── aggregate_table ──────────────────────────────────────────────
    tool({
      name: "aggregate_table",
      description: "Perform aggregation operations on a numeric column",
      parameters: {
        database: z.string().min(1),
        table: z.string().min(1),
        column: z.string().min(1).describe("Column to aggregate"),
        operation: z
          .enum(["count", "sum", "avg", "min", "max"])
          .describe("Aggregation operation"),
      },
      implementation: async ({ database, table, column, operation }) => {
        const dbDir = `${dataRoot}/databases/${database}`;
        const allRows = await csvEngine.parseCSV(`${dbDir}/${table}.csv`);
        const values = allRows
          .map((r) => Number(r[column]))
          .filter((v) => !isNaN(v));
        const count = values.length;

        let result: number;
        switch (operation) {
          case "count":
            result = count;
            break;
          case "sum":
            result = values.reduce((a, b) => a + b, 0);
            break;
          case "avg":
            result = count > 0 ? values.reduce((a, b) => a + b, 0) / count : 0;
            break;
          case "min":
            result = count > 0 ? Math.min(...values) : 0;
            break;
          case "max":
            result = count > 0 ? Math.max(...values) : 0;
            break;
        }

        return JSON.stringify({ operation, column, result, count });
      },
    }),
  ];
}

async function fsRm(fp: string): Promise<void> {
  try {
    const fs = await import("node:fs/promises");
    await fs.rm(fp, { force: true });
  } catch {
    // ignore
  }
}