throughline

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.js";

const ENCODING = "utf-8" as const;
const COMPRESSION_THRESHOLD = 0.65;

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

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

function getTimestamp(): string {
  return new Date().toISOString().replace(/[:.]/g, "-").slice(0, 19);
}

async function fetchJson(url: string, apiToken: string): Promise<unknown> {
  const response = await fetch(url, {
    headers: apiToken ? { "Authorization": `Bearer ${apiToken}` } : {},
  });
  if (!response.ok) throw new Error(`HTTP ${response.status} from ${url}`);
  return response.json();
}

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

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

  function getConfig() {
    const cfg = ctl.getPluginConfig(configSchematics);
    return {
      apiBase: cfg.get("apiBase").trim() || "http://127.0.0.1:1234",
      apiToken: cfg.get("apiToken").trim(),
      handoffDir: cfg.get("handoffDir").trim() || "C:\\Users\\Kevin\\.lmstudio\\handoff",
    };
  }

  // ── check_context ─────────────────────────────────────────────────────────

  const checkContextTool = tool({
    name: "check_context",
    description:
      "Checks current session token usage against the model's context window. " +
      "Call this at the start of every response. " +
      "If the result indicates compression is needed, immediately call compress_session without finishing this response.",
    parameters: {
      model_id: z.string().describe("The currently loaded model ID."),
      previous_response_id: z.string().describe("The response_id from the previous turn. Empty string if first turn."),
      prompt_tokens: z.number().describe("The prompt_tokens value from the previous API response. 0 if first turn."),
    },
    implementation: async ({ model_id, previous_response_id, prompt_tokens }) => {
      const { apiBase, apiToken } = getConfig();

      try {
        // Get model's max context length
        const modelsData = await fetchJson(`${apiBase}/api/v1/models`, apiToken) as {
          models?: Array<{ key: string; max_context_length?: number }>;
        };

        const model = modelsData.models?.find(m => m.key === model_id);
        const maxContext = model?.max_context_length;

        if (!maxContext) {
          return `Unable to determine context length for model '${model_id}'. Skipping check.`;
        }

        if (prompt_tokens === 0) {
          return `Session start. Context window: ${maxContext} tokens. No usage data yet.`;
        }

        const usage = prompt_tokens / maxContext;
        const pct = (usage * 100).toFixed(1);

        if (usage >= COMPRESSION_THRESHOLD) {
          return `COMPRESS NOW: ${pct}% of context used (${prompt_tokens}/${maxContext} tokens). Call compress_session immediately.`;
        }

        return `Context OK: ${pct}% used (${prompt_tokens}/${maxContext} tokens).`;
      } catch (err) {
        return `Context check failed: ${String(err)}`;
      }
    },
  });

  // ── compress_session ──────────────────────────────────────────────────────

  const compressSessionTool = tool({
    name: "compress_session",
    description:
      "Compresses the current session into a structured handoff document and writes it to disk. " +
      "Call this immediately when check_context returns COMPRESS NOW, or when you detect any of these self-report signals: " +
      "(1) uncertainty about something established earlier in the session, " +
      "(2) hedging on something you should know confidently, " +
      "(3) significant shift in conversation focus since the session began. " +
      "\n\n" +
      "When called, produce a handoff document with exactly these three sections:\n\n" +
      "[PERMANENT CORE]\n" +
      "Facts that would materially change behavior if missing. The session purpose, key decisions never to be revisited, foundational constraints. " +
      "Ruthlessly lean — if its absence would not break the next session, cut it. Never pruned.\n\n" +
      "[CURRENT STATE]\n" +
      "Where things stand right now. Active protocols, recent decisions, open questions, next steps. " +
      "Terse bullet points only. Rewritten fresh each compression pass.\n\n" +
      "[RECENT CONVERSATION]\n" +
      "Last 4 exchanges verbatim. User messages and assistant responses both. No summarizing — exact text.\n\n" +
      "Rules: Every Permanent Core fact must pass this test: would its absence materially change the next session? " +
      "Output only the three sections. No preamble, no closing remarks, no meta-commentary.",
    parameters: {
      handoff_document: z.string().describe(
        "The complete handoff document with all three sections: [PERMANENT CORE], [CURRENT STATE], and [RECENT CONVERSATION]."
      ),
    },
    implementation: async ({ handoff_document }) => {
      const { handoffDir } = getConfig();

      try {
        const timestamp = getTimestamp();
        const latestPath = `${handoffDir}\\latest.md`;
        const archivePath = `${handoffDir}\\archive\\handoff-${timestamp}.md`;

        await ensureDir(latestPath);
        await ensureDir(archivePath);

        const content = `# throughLine Handoff\nGenerated: ${new Date().toISOString()}\n\n${handoff_document.trim()}\n`;

        await writeFile(latestPath, content, ENCODING);
        await writeFile(archivePath, content, ENCODING);

        return (
          `Compression complete.\n\n` +
          `Start a new session and attach:\n` +
          `${latestPath}\n\n` +
          `Archive saved to:\n` +
          `${archivePath}`
        );
      } catch (err) {
        return `Compression failed: ${String(err)}`;
      }
    },
  });

  // ─────────────────────────────────────────────────────────────────────────

  return [
    checkContextTool,
    compressSessionTool,
  ];
}