simplememory

import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { existsSync } from "node:fs";
import { mkdir, readFile, writeFile } from "node:fs/promises";
import { dirname } from "node:path";
import { z } from "zod";
import { configSchematics } from "./config";

const SIZE_WARNING_CHARS = 8000;
const ENCODING = "utf-8" as const;

// ─── Helpers ────────────────────────────────────────────────────────────────

async function ensureMemoryFile(filePath: string): Promise<void> {
  const dir = dirname(filePath);
  if (!existsSync(dir)) await mkdir(dir, { recursive: true });
  if (!existsSync(filePath)) await writeFile(filePath, "", ENCODING);
}

async function resolveFile(getPath: () => string): Promise<{ filePath: string; content: string }> {
  const filePath = getPath();
  await ensureMemoryFile(filePath);
  const content = normalizeNewlines(await readFile(filePath, ENCODING));
  return { filePath, content };
}

function getTimestamp(): string {
  return new Date().toISOString().split("T")[0];
}

function normalizeNewlines(text: string): string {
  return text.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
}

// Strip backslash escapes added by markdown renderers (e.g. \& -> &, \~ -> ~, \_ -> _)
function stripMarkdownEscapes(text: string): string {
  return text.replace(/\\([&~_*#`!|{}\[\]().])/g, "$1");
}

function countEntries(content: string): number {
  return (content.match(/^\[20\d\d-\d\d-\d\d\]/gm) || []).length;
}

function buildSizeWarning(content: string): string {
  if (content.length <= SIZE_WARNING_CHARS) return "";
  const entries = countEntries(content);
  const dates = [...content.matchAll(/^\[(\d{4}-\d{2}-\d{2})\]/gm)].map(m => m[1]);
  const uniqueDays = new Set(dates).size;
  const density = uniqueDays > 0 ? (entries / uniqueDays).toFixed(1) : "N/A";
  return `\n[Memory warning: ${content.length} chars, ${entries} entries, ${density} entries/day average. Consolidate before responding.]`;
}

function findSection(lines: string[], sectionName: string): { start: number; end: number } | null {
  const header = stripMarkdownEscapes(`## ${sectionName}`);
  const startIdx = lines.findIndex(l => stripMarkdownEscapes(l.trim()) === header);
  if (startIdx === -1) return null;
  let endIdx = lines.length;
  for (let i = startIdx + 1; i < lines.length; i++) {
    if (lines[i].startsWith("## ") || stripMarkdownEscapes(lines[i]).startsWith("## ")) { endIdx = i; break; }
  }
  return { start: startIdx, end: endIdx };
}

function matchAndReplace(content: string, oldText: string, newText: string): string | null {
  const normalizedOld = stripMarkdownEscapes(normalizeNewlines(oldText.trim()));
  const normalizedNew = normalizeNewlines(newText.trim());
  const normalizedContent = stripMarkdownEscapes(content);

  if (normalizedContent.includes(normalizedOld)) {
    return normalizedContent.replace(normalizedOld, normalizedNew).replace(/\n{3,}/g, "\n\n");
  }

  // Fallback: try against original content without escape normalization
  if (content.includes(normalizedOld)) {
    return content.replace(normalizedOld, normalizedNew).replace(/\n{3,}/g, "\n\n");
  }

  const escaped = normalizedOld.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
  const fuzzy = new RegExp(escaped.replace(/\s+/g, "\\s+"), "");
  if (fuzzy.test(normalizedContent)) {
    return normalizedContent.replace(fuzzy, normalizedNew).replace(/\n{3,}/g, "\n\n");
  }

  return null;
}

// ─── Tool Provider ───────────────────────────────────────────────────────────

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {

  function getMemoryPath(): string {
    const configured = ctl.getPluginConfig(configSchematics).get("memoryPath").trim();
    if (!configured) throw new Error("Memory file path is not configured. Please set it in the plugin settings.");
    return configured;
  }

  const getCurrentTimeTool = tool({
    name: "get_current_time",
    description: "Returns the current local date and time from the system clock. Call this at the start of every session after read_memory.",
    parameters: {},
    implementation: async () => {
      const now = new Date();
      const date = now.toLocaleDateString("en-US", { weekday: "long", year: "numeric", month: "long", day: "numeric" });
      const time = now.toLocaleTimeString("en-US", { hour: "2-digit", minute: "2-digit", timeZoneName: "short" });
      return `${date} ${time}`;
    },
  });

  const readMemoryTool = tool({
    name: "read_memory",
    description:
      "Reads the entire persistent memory file. Call this at the start of every session before responding to anything. " +
      "If the response includes a memory warning, consolidate before proceeding.",
    parameters: {},
    implementation: async () => {
      const { content } = await resolveFile(getMemoryPath);
      if (!content.trim()) return "Memory is empty. No prior context available.";
      return content + buildSizeWarning(content);
    },
  });

  const saveMemoryTool = tool({
    name: "save_memory",
    description:
      "Appends a new dated entry to the memory file. Use for standalone facts not belonging to a structured section. " +
      "Write densely: subject, verb, object. Third person, present tense. No filler phrases, no 'user mentioned.' " +
      "One fact or tightly related cluster per call. Skip if similar content already exists — use update_memory instead.",
    parameters: {
      fact: z.string().describe(
        "The fact to save. Dense, third-person, present tense. " +
        "Example: 'Kevin uses liquid melatonin 0.3-0.6mg for phase-shifting, not sedation.'"
      ),
    },
    implementation: async ({ fact }) => {
      const { filePath, content } = await resolveFile(getMemoryPath);
      const normalizedFact = normalizeNewlines(fact.trim());
      const preview = normalizedFact.slice(0, 60).toLowerCase();
      if (preview.length > 20 && stripMarkdownEscapes(content).toLowerCase().includes(preview)) {
        return "Skipped: similar content already exists in memory. Use update_memory to correct it.";
      }
      await writeFile(filePath, content + `[${getTimestamp()}] ${normalizedFact}\n`, ENCODING);
      return "Saved.";
    },
  });

  const updateMemoryTool = tool({
    name: "update_memory",
    description:
      "Finds and replaces an existing passage in memory. Use instead of save_memory when correcting or updating existing information. " +
      "Pass an empty string for replace to delete the entry. For structured sections, prefer replace_section. Can also rename section headers by finding and replacing the ## header line directly.",
    parameters: {
      find: z.string().describe("The text to find in memory."),
      replace: z.string().describe("Replacement text. Empty string to delete."),
    },
    implementation: async ({ find, replace }) => {
      const { filePath, content } = await resolveFile(getMemoryPath);
      const updated = matchAndReplace(content, find, replace);
      if (updated === null) return "Could not find the specified text in memory. No changes made.";
      await writeFile(filePath, updated, ENCODING);
      const deleted = replace.trim() === "";
      return deleted ? "Deleted." : "Updated.";
    },
  });

  const deleteMemoryTool = tool({
    name: "delete_memory",
    description:
      "Removes a specific passage from memory entirely. Use when asked to forget something or remove an entry. " +
      "More reliable than update_memory for multi-line entries.",
    parameters: {
      entry: z.string().describe("The exact text to remove from memory."),
    },
    implementation: async ({ entry }) => {
      const { filePath, content } = await resolveFile(getMemoryPath);
      const updated = matchAndReplace(content, entry, "");
      if (updated === null) return "Could not find the specified text in memory. No changes made.";
      await writeFile(filePath, updated, ENCODING);
      return "Deleted.";
    },
  });

  const searchMemoryTool = tool({
    name: "search_memory",
    description:
      "Returns only the entries in memory containing the query string. " +
      "More token-efficient than read_memory when looking for a specific detail.",
    parameters: {
      query: z.string().describe("The word or phrase to search for in memory."),
    },
    implementation: async ({ query }) => {
      const { content } = await resolveFile(getMemoryPath);
      const lower = query.toLowerCase().trim();
      if (!lower) return "Query must not be empty.";
      const lines = content.split("\n");
      const matches: string[] = [];
      for (let i = 0; i < lines.length; i++) {
        if (lines[i].toLowerCase().includes(lower)) {
          const block = lines.slice(i, Math.min(i + 3, lines.length)).join("\n").trim();
          if (block && !matches.includes(block)) matches.push(block);
        }
      }
      if (matches.length === 0) return `No entries found containing "${query}".`;
      return `Found ${matches.length} match(es):\n\n${matches.join("\n\n")}`;
    },
  });

  const saveToSectionTool = tool({
    name: "save_to_section",
    description:
      "Appends content inside a named section of the memory file. " +
      "Use for structured categories: 'Health & Wellness', 'Hardware & Tech Setup', 'Home Inventory', 'Preferences & Notes'. " +
      "Creates the section if it does not exist. Use bullet points for lists. Dense, present tense, no filler.",
    parameters: {
      target: z.string().describe(
        "Section name without ##. Canonical sections: 'Health & Wellness', 'Hardware & Tech Setup', 'Home Inventory', 'Preferences & Notes'."
      ),
      content: z.string().describe(
        "Content to add. Use - bullet points (not *). Dense, present tense. Example: '- RTX 5090 (32GB VRAM)'"
      ),
    },
    implementation: async ({ target, content }) => {
      const { filePath, content: raw } = await resolveFile(getMemoryPath);
      const lines = raw.split("\n");
      const sectionRange = findSection(lines, target);
      const normalizedContent = normalizeNewlines(content.trim());

      if (sectionRange === null) {
        await writeFile(filePath, raw.trimEnd() + `\n## ${target}\n\n${normalizedContent}\n`, ENCODING);
        return `Section '${target}' created and content added.`;
      }

      let insertLine = sectionRange.start + 1;
      for (let i = sectionRange.end - 1; i > sectionRange.start; i--) {
        if (lines[i].trim() !== "") { insertLine = i + 1; break; }
      }

      const newLines = [...lines.slice(0, insertLine), normalizedContent, ...lines.slice(insertLine)];
      await writeFile(filePath, newLines.join("\n"), ENCODING);
      return `Added to section '${target}'.`;
    },
  });

  const replaceSectionTool = tool({
    name: "replace_section",
    description:
      "Replaces the entire body of a named section in the memory file. " +
      "Use for wholesale updates to structured sections like hardware specs or pantry inventory. " +
      "More reliable than update_memory for multi-line section content. Preserves the section header.",
    parameters: {
      target: z.string().describe(
        "Section name without ##. Must match an existing header exactly. " +
        "Examples: 'Health & Wellness', 'Hardware & Tech Setup', 'Home Inventory'."
      ),
      new_content: z.string().describe(
        "Full replacement content for the section body. Use - bullet points (not *). Dense, present tense, no filler."
      ),
    },
    implementation: async ({ target, new_content }) => {
      const { filePath, content } = await resolveFile(getMemoryPath);
      const lines = content.split("\n");
      const sectionRange = findSection(lines, target);

      if (sectionRange === null) {
        return `Section '${target}' not found in memory. Use save_to_section to create it.`;
      }

      const normalizedContent = normalizeNewlines(new_content.trim());
      const newLines = [
        ...lines.slice(0, sectionRange.start),
        lines[sectionRange.start], "",
        normalizedContent, "",
        ...lines.slice(sectionRange.end),
      ].join("\n").replace(/\n{3,}/g, "\n\n");

      await writeFile(filePath, newLines, ENCODING);
      return `Section '${target}' replaced.`;
    },
  });

  return [
    getCurrentTimeTool,
    readMemoryTool,
    saveMemoryTool,
    updateMemoryTool,
    deleteMemoryTool,
    searchMemoryTool,
    saveToSectionTool,
    replaceSectionTool,
  ];
}