knowledge

import { text, tool, type Tool, type ToolsProvider } from "@lmstudio/sdk";
import { z } from "zod";
import { pluginConfigSchematics } from "./config";
import {
  getDataDir, getDb,
  insertNote, updateNote, deleteNote, getNoteById, listNotes, searchNotes,
  upsertConcept, getConceptByName, listConcepts,
  linkNoteToConcept, unlinkNoteConcept,
  getConceptNotes, getNoteConcepts, getUnlinkedNotes, countNotes,
} from "./db";

function json(obj: unknown): string {
  return JSON.stringify(obj, null, 2);
}

function safe_impl<T extends Record<string, unknown>>(
  name: string,
  fn: (params: T) => Promise<string>
): (params: T) => Promise<string> {
  return async (params: T) => {
    try {
      return await fn(params);
    } catch (err: unknown) {
      const msg = err instanceof Error ? err.message : String(err);
      return JSON.stringify({
        tool_error: true,
        tool: name,
        error: msg,
        hint: "Read the error above, fix the parameter causing the issue, and retry the tool call.",
      }, null, 2);
    }
  };
}

const FETCH_HEADERS = {
  "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
  "Accept": "text/html,application/xhtml+xml,*/*;q=0.8",
  "Accept-Language": "en-US,en;q=0.9",
};

function extractText(html: string, maxChars = 12000): string {
  const articleMatch =
    html.match(/<article[^>]*>([\s\S]*?)<\/article>/i) ??
    html.match(/<main[^>]*>([\s\S]*?)<\/main>/i);
  const src = articleMatch ? articleMatch[1] : html;
  return src
    .replace(/<script[\s\S]*?<\/script>/gi, " ")
    .replace(/<style[\s\S]*?<\/style>/gi, " ")
    .replace(/<nav[\s\S]*?<\/nav>/gi, " ")
    .replace(/<[^>]+>/g, " ")
    .replace(/&nbsp;/g, " ").replace(/&amp;/g, "&").replace(/&lt;/g, "<").replace(/&gt;/g, ">")
    .replace(/\s{2,}/g, " ")
    .trim()
    .slice(0, maxChars);
}

function extractTitle(html: string): string {
  const m = html.match(/<title[^>]*>([^<]{1,200})<\/title>/i);
  if (m) return m[1].trim().replace(/\s+/g, " ");
  return "Untitled";
}

async function fetchUrl(url: string, timeoutMs = 15_000): Promise<{ title: string; text: string }> {
  const res = await fetch(url, { signal: AbortSignal.timeout(timeoutMs), headers: FETCH_HEADERS });
  if (!res.ok) throw new Error(`HTTP ${res.status}`);
  const html = await res.text();
  return { title: extractTitle(html), text: extractText(html) };
}

export const toolsProvider: ToolsProvider = async (ctl) => {
  const cfg = ctl.getPluginConfig(pluginConfigSchematics);
  const db = () => getDb(getDataDir(cfg.get("dataPath")));

  const tools: Tool[] = [

    // =========================================================================
    // NOTE MANAGEMENT
    // =========================================================================

    tool({
      name: "ingest_note",
      description: text`
        Save a note to the knowledge base with title, content, source, and tags.
        Returns the saved note ID.
        Call this whenever the user shares text, a quote, an idea, or a summary
        they want to remember. Do NOT call for ephemeral conversation context.
      `,
      parameters: {
        title: z.string().min(3).describe("Short descriptive title for the note"),
        content: z.string().min(10).describe("Full text content of the note"),
        source: z.string().default("").describe("Where this came from — URL, book title, person name, etc."),
        tags: z.array(z.string()).default([]).describe("Tags for filtering and concept association"),
      },
      implementation: safe_impl("ingest_note", async ({ title, content, source, tags }) => {
        const id = insertNote(db(), title, content, source, tags);
        return json({ success: true, id, title, tags, source });
      }),
    }),

    tool({
      name: "ingest_url",
      description: text`
        Fetch a URL, extract the page text, and save it as a note.
        Returns the note ID and extracted text preview.
        Use when the user shares a link they want to add to their knowledge base.
        Do NOT use for general web search — use web-search tools for that.
      `,
      parameters: {
        url: z.string().url().describe("URL to fetch and ingest"),
        tags: z.array(z.string()).default([]).describe("Tags to apply to the ingested note"),
        overrideTitle: z.string().default("").describe("Optional title override — if blank, the page title is used"),
      },
      implementation: safe_impl("ingest_url", async ({ url, tags, overrideTitle }) => {
        const { title: pageTitle, text: content } = await fetchUrl(url);
        const title = overrideTitle.trim() || pageTitle;
        const id = insertNote(db(), title, content, url, tags);
        return json({
          success: true,
          id,
          title,
          source: url,
          tags,
          contentPreview: content.slice(0, 300) + (content.length > 300 ? "…" : ""),
          wordCount: content.split(/\s+/).length,
        });
      }),
    }),

    tool({
      name: "update_note",
      description: text`
        Update fields on an existing note by ID.
        Only supply fields to change. Supply tags as a full replacement array, not a delta.
        Do NOT call if the note was just ingested — it was already saved correctly.
      `,
      parameters: {
        id: z.coerce.number().int().describe("Note ID"),
        title: z.string().optional(),
        content: z.string().optional(),
        source: z.string().optional(),
        tags: z.array(z.string()).optional(),
      },
      implementation: safe_impl("update_note", async ({ id, ...patch }) => {
        updateNote(db(), id, patch);
        const updated = getNoteById(db(), id);
        return json({ success: true, note: { ...updated, tags: JSON.parse(updated.tags) } });
      }),
    }),

    tool({
      name: "delete_note",
      description: text`
        Permanently delete a note by ID. Also removes all concept links for this note.
        Confirm the ID before calling — deletion is irreversible.
      `,
      parameters: {
        id: z.coerce.number().int().describe("Note ID to delete"),
      },
      implementation: safe_impl("delete_note", async ({ id }) => {
        deleteNote(db(), id);
        return json({ success: true, deleted: id });
      }),
    }),

    tool({
      name: "get_note",
      description: text`
        Get the full content of a note by ID.
        Returns title, content, source, tags, and linked concepts.
        Call before update_note to confirm you have the right note.
      `,
      parameters: {
        id: z.coerce.number().int().describe("Note ID"),
      },
      implementation: safe_impl("get_note", async ({ id }) => {
        const note = getNoteById(db(), id);
        const concepts = getNoteConcepts(db(), id);
        return json({
          ...note,
          tags: JSON.parse(note.tags),
          linkedConcepts: concepts.map((c) => ({ id: c.id, name: c.name })),
        });
      }),
    }),

    tool({
      name: "list_notes",
      description: text`
        List notes with optional filtering by tag or date.
        Returns title, source, tags, createdAt — no full content (use get_note for that).
        Use before get_note or update_note to find the right note ID.
        Do NOT call if you already have the note ID.
      `,
      parameters: {
        tag: z.string().default("").describe("Filter to notes with this tag"),
        since: z.string().default("").describe("ISO date — only notes created on or after this date"),
        limit: z.coerce.number().int().min(1).max(200).default(50).describe("Max results"),
      },
      implementation: safe_impl("list_notes", async ({ tag, since, limit }) => {
        const notes = listNotes(db(), {
          tag: tag || undefined,
          since: since || undefined,
          limit,
        });
        return json({
          total: notes.length,
          notes: notes.map((n) => ({
            id: n.id,
            title: n.title,
            source: n.source,
            tags: JSON.parse(n.tags),
            createdAt: n.createdAt.slice(0, 10),
          })),
        });
      }),
    }),

    tool({
      name: "search_notes",
      description: text`
        Full-text search across all notes — searches title, content, source, and tags.
        Returns matching notes ranked by relevance.
        Call this before generate_synthesis to gather relevant notes on a topic.
        Do NOT call for exact ID lookups — use get_note for that.
      `,
      parameters: {
        query: z.string().min(2).describe("Search query — keywords, phrases, or concepts"),
        limit: z.coerce.number().int().min(1).max(50).default(15).describe("Max results"),
      },
      implementation: safe_impl("search_notes", async ({ query, limit }) => {
        const notes = searchNotes(db(), query, limit);
        return json({
          query,
          total: notes.length,
          notes: notes.map((n) => ({
            id: n.id,
            title: n.title,
            source: n.source,
            tags: JSON.parse(n.tags),
            createdAt: n.createdAt.slice(0, 10),
            excerpt: n.content.slice(0, 200) + (n.content.length > 200 ? "…" : ""),
          })),
        });
      }),
    }),

    // =========================================================================
    // CONCEPT GRAPH
    // =========================================================================

    tool({
      name: "upsert_concept",
      description: text`
        Create a new concept or update an existing one's description.
        Concepts are the organizing nodes of the knowledge graph — themes, topics, people, frameworks.
        Returns the concept ID.
        Call after ingesting a cluster of related notes to give them a shared concept home.
      `,
      parameters: {
        name: z.string().min(2).describe("Concept name — unique, short, capitalized"),
        description: z.string().default("").describe("What this concept covers and why it matters"),
      },
      implementation: safe_impl("upsert_concept", async ({ name, description }) => {
        const concept = upsertConcept(db(), name, description);
        return json({ success: true, concept });
      }),
    }),

    tool({
      name: "link_note_to_concept",
      description: text`
        Associate a note with a concept in the knowledge graph.
        A note can be linked to multiple concepts. Idempotent — safe to call twice.
        Call after upsert_concept to build the graph connections.
        Do NOT call without first confirming both the note ID and concept exist.
      `,
      parameters: {
        noteId: z.coerce.number().int().describe("Note ID"),
        conceptName: z.string().describe("Concept name — must already exist (call upsert_concept first if needed)"),
      },
      implementation: safe_impl("link_note_to_concept", async ({ noteId, conceptName }) => {
        getNoteById(db(), noteId);
        let concept = getConceptByName(db(), conceptName);
        if (!concept) {
          concept = upsertConcept(db(), conceptName, "");
        }
        linkNoteToConcept(db(), noteId, concept.id);
        return json({ success: true, noteId, conceptId: concept.id, conceptName });
      }),
    }),

    tool({
      name: "unlink_note_from_concept",
      description: text`
        Remove the association between a note and a concept.
        Does not delete the note or the concept — only removes the link.
        Call when a note was incorrectly linked.
      `,
      parameters: {
        noteId: z.coerce.number().int().describe("Note ID"),
        conceptName: z.string().describe("Concept name"),
      },
      implementation: safe_impl("unlink_note_from_concept", async ({ noteId, conceptName }) => {
        const concept = getConceptByName(db(), conceptName);
        if (!concept) throw new Error(`Concept '${conceptName}' not found.`);
        unlinkNoteConcept(db(), noteId, concept.id);
        return json({ success: true, noteId, conceptName });
      }),
    }),

    tool({
      name: "get_concept",
      description: text`
        Get a concept by name with all its linked notes.
        Returns the concept description and full list of notes associated with it.
        Call before generate_synthesis to load the relevant note set for a concept.
      `,
      parameters: {
        name: z.string().describe("Concept name"),
      },
      implementation: safe_impl("get_concept", async ({ name }) => {
        const concept = getConceptByName(db(), name);
        if (!concept) throw new Error(`Concept '${name}' not found.`);
        const notes = getConceptNotes(db(), concept.id);
        return json({
          concept,
          noteCount: notes.length,
          notes: notes.map((n) => ({
            id: n.id,
            title: n.title,
            source: n.source,
            tags: JSON.parse(n.tags),
            createdAt: n.createdAt.slice(0, 10),
            excerpt: n.content.slice(0, 300) + (n.content.length > 300 ? "…" : ""),
          })),
        });
      }),
    }),

    tool({
      name: "list_concepts",
      description: text`
        List all concepts in the knowledge graph, sorted by note count.
        Returns concept name, description, and note count.
        Use this for a map of everything you've organized so far.
        Call at session start to orient to the current knowledge base.
      `,
      parameters: {},
      implementation: safe_impl("list_concepts", async () => {
        const concepts = listConcepts(db());
        return json({
          total: concepts.length,
          concepts: concepts.map((c) => ({
            id: c.id,
            name: c.name,
            description: c.description,
            noteCount: c.noteCount,
          })),
        });
      }),
    }),

    // =========================================================================
    // SYNTHESIS & REVIEW
    // =========================================================================

    tool({
      name: "generate_synthesis",
      description: text`
        Scaffold a synthesis of notes on a concept or topic.
        Loads notes (from a concept or search), then returns a structured prompt
        with those notes for the LLM to synthesize into a coherent summary,
        identify recurring themes, surface contradictions, and note gaps.
        Returns the synthesis payload — no side effects.
        Call after get_concept or search_notes to synthesize what you've captured.
      `,
      parameters: {
        conceptName: z.string().default("").describe("Concept name to synthesize. Leave blank to use noteIds or query instead."),
        noteIds: z.array(z.coerce.number().int()).default([]).describe("Specific note IDs to include in synthesis"),
        searchQuery: z.string().default("").describe("Search query to gather notes for synthesis. Used if conceptName and noteIds are both blank."),
        focusQuestion: z.string().default("").describe("Optional: a specific question the synthesis should answer"),
      },
      implementation: safe_impl("generate_synthesis", async ({ conceptName, noteIds, searchQuery, focusQuestion }) => {
        let notes: Array<{ id: number; title: string; source: string; content: string }> = [];

        if (conceptName) {
          const concept = getConceptByName(db(), conceptName);
          if (!concept) throw new Error(`Concept '${conceptName}' not found.`);
          const rows = getConceptNotes(db(), concept.id);
          notes = rows.map((n) => ({ id: n.id, title: n.title, source: n.source, content: n.content }));
        } else if (noteIds.length > 0) {
          notes = noteIds.map((id) => {
            const n = getNoteById(db(), id);
            return { id: n.id, title: n.title, source: n.source, content: n.content };
          });
        } else if (searchQuery) {
          const rows = searchNotes(db(), searchQuery, 15);
          notes = rows.map((n) => ({ id: n.id, title: n.title, source: n.source, content: n.content }));
        } else {
          throw new Error("Provide one of: conceptName, noteIds, or searchQuery.");
        }

        if (notes.length === 0) {
          return json({ action: "no_notes", message: "No notes found. Ingest some notes first." });
        }

        return json({
          action: "synthesize",
          concept: conceptName || null,
          focusQuestion: focusQuestion || null,
          noteCount: notes.length,
          notes,
          instructions: [
            focusQuestion
              ? `Answer this question using the notes above: "${focusQuestion}"`
              : `Synthesize the notes above into a coherent summary${conceptName ? ` on the concept "${conceptName}"` : ""}.`,
            "Structure your synthesis:",
            "1. CORE THESIS — what is the central idea or conclusion across these notes?",
            "2. KEY THEMES — 3–5 recurring themes with evidence from specific notes (cite by title)",
            "3. SUPPORTING DETAILS — facts, quotes, or examples that strengthen the core thesis",
            "4. CONTRADICTIONS — where notes conflict or tension exists — do not resolve artificially",
            "5. GAPS — what is NOT covered that should be? What follow-up questions arise?",
            "6. CONNECTIONS — unexpected links between notes from different sources",
            "Be specific: cite note titles and sources. Never generalize beyond what the notes say.",
          ].join("\n"),
        });
      }),
    }),

    tool({
      name: "knowledge_review",
      description: text`
        Weekly knowledge base review — surfaces what's been added recently,
        which notes are not yet linked to any concept (orphans),
        which concepts have the most growth, and where gaps exist.
        Use at the start of a session or after a burst of ingestion to organize.
      `,
      parameters: {
        daysSince: z.coerce.number().int().min(1).max(365).default(7)
          .describe("Review notes added in the last N days"),
      },
      implementation: safe_impl("knowledge_review", async ({ daysSince }) => {
        const since = new Date(Date.now() - daysSince * 86_400_000).toISOString();
        const recentNotes = listNotes(db(), { since, limit: 50 });
        const orphans = getUnlinkedNotes(db(), 20);
        const concepts = listConcepts(db());

        const totalNotes = countNotes(db());

        const topConcepts = concepts.slice(0, 10).map((c) => ({
          name: c.name,
          noteCount: c.noteCount,
          description: c.description.slice(0, 100),
        }));

        const emptyConcepts = concepts.filter((c) => c.noteCount === 0);

        return json({
          summary: {
            totalNotes,
            totalConcepts: concepts.length,
            recentNotesCount: recentNotes.length,
            orphanNotesCount: orphans.length,
            emptyConcepts: emptyConcepts.length,
          },
          recentNotes: recentNotes.map((n) => ({
            id: n.id,
            title: n.title,
            source: n.source,
            tags: JSON.parse(n.tags),
            createdAt: n.createdAt.slice(0, 10),
          })),
          orphanNotes: orphans.map((n) => ({
            id: n.id,
            title: n.title,
            createdAt: n.createdAt.slice(0, 10),
          })),
          topConcepts,
          emptyConcepts: emptyConcepts.map((c) => c.name),
          suggestions: [
            orphans.length > 0 ? `Link ${orphans.length} unlinked note(s) to concepts.` : null,
            emptyConcepts.length > 0 ? `${emptyConcepts.length} concept(s) have no linked notes — add notes or delete them.` : null,
            recentNotes.length > 0 ? `Review ${recentNotes.length} recent note(s) and synthesize any clusters.` : null,
          ].filter(Boolean),
        });
      }),
    }),

  ];

  return tools;
};