import { Chat, type ChatMessage, type PromptPreprocessorController } from "@lmstudio/sdk";
import { clamp, stripReasoning } from "./workspace";

/**
 * Automatic context compaction.
 *
 * A plugin cannot delete earlier turns -- a prompt preprocessor may only rewrite
 * the newest user message. What it CAN do is notice that the conversation is
 * approaching the model's context limit and attach a running summary of the
 * older turns to the newest message. When the host then drops the oldest
 * messages to make room, the summary survives, because it now lives at the end
 * of the conversation rather than the beginning.
 *
 * The net effect matches auto-compaction: work done thirty turns ago is still
 * known to the model after the raw turns have fallen out of the window.
 */

const OPEN_TAG = "<earlier-conversation-summary>";
const CLOSE_TAG = "</earlier-conversation-summary>";

/** Turns kept verbatim at the end; only what precedes them is summarised. */
const KEEP_RECENT = 6;
/** Re-summarise only after the conversation has grown by this many messages. */
const RESUMMARISE_EVERY = 8;
const SUMMARY_TOKENS = 600;
const MAX_SUMMARY_CHARS = 4000;
const MAX_SOURCE_CHARS = 60000;
const MAX_TRACKED_CHATS = 12;

/** The subset of an LLM handle compaction needs, feature-detected at runtime. */
interface MeasurableModel {
  getContextLength?: () => Promise<number>;
  countTokens?: (input: string) => Promise<number>;
  respond: (chat: Chat, opts: { maxTokens: number; temperature: number }) => unknown;
}

interface CompactionState {
  summary: string;
  /** Number of messages already folded into the summary. */
  covered: number;
  /** Guards against two overlapping summarisations in one conversation. */
  busy: boolean;
}

const states = new Map<string, CompactionState>();

/**
 * Conversations carry no id, so identify one by its opening exchange. Stable for
 * the life of a chat, and different chats effectively never collide.
 */
function fingerprint(messages: ChatMessage[]): string {
  // Three messages rather than two: chats often open identically ("make a
  // platformer"), and a collision would leak one chat's summary into another
  // until the next refresh.
  const opening = messages.slice(0, 3).map(readText);
  const shape = opening.map((text) => text.length).join(",");
  return `${shape}|${opening.join("~").slice(0, 300)}`;
}

function readText(message: unknown): string {
  if (typeof message === "string") return message;
  if (message === null || typeof message !== "object") return "";
  const candidate = message as { getText?: () => unknown; content?: unknown; text?: unknown };
  try {
    if (typeof candidate.getText === "function") {
      const value = candidate.getText();
      if (typeof value === "string") return value;
    }
  } catch {
    // Fall through.
  }
  if (typeof candidate.text === "string") return candidate.text;
  if (typeof candidate.content === "string") return candidate.content;
  return "";
}

function readRole(message: unknown): string {
  const candidate = message as { getRole?: () => unknown; role?: unknown };
  try {
    if (typeof candidate?.getRole === "function") {
      const value = candidate.getRole();
      if (typeof value === "string") return value;
    }
  } catch {
    // Fall through.
  }
  return typeof candidate?.role === "string" ? candidate.role : "unknown";
}

export interface CompactionResult {
  /** Block to inject, or "" when there is nothing to add. */
  block: string;
  /** Diagnostic line for the plugin log. */
  note: string;
}

/**
 * Returns a summary block when the conversation is close to filling the context
 * window. Never throws: on any failure the caller simply gets no block.
 */
export async function maybeCompact(
  ctl: PromptPreprocessorController,
  triggerPercent: number,
): Promise<CompactionResult> {
  try {
    const history = await ctl.pullHistory();
    const messages = history.getMessagesArray();
    if (messages.length < KEEP_RECENT + 2) return { block: "", note: "" };

    const key = fingerprint(messages);
    const existing = states.get(key);

    // tokenSource() may hand back a generator handle, which cannot measure
    // tokens. Feature-detect rather than assume a full LLM.
    const model = (await ctl.tokenSource()) as unknown as MeasurableModel;
    if (typeof model.getContextLength !== "function" || typeof model.countTokens !== "function") {
      return { block: "", note: "" };
    }
    const contextLength = await model.getContextLength();
    if (!Number.isFinite(contextLength) || contextLength <= 0) return { block: "", note: "" };

    const transcript = messages.map((m) => `${readRole(m)}: ${readText(m)}`).join("\n");
    const used = await model.countTokens(transcript);
    const percent = Math.round((used / contextLength) * 100);

    // Below the trigger: keep serving any summary already built, so the model
    // does not lose it again the moment usage dips.
    if (percent < triggerPercent) {
      return {
        block: existing === undefined ? "" : render(existing.summary, percent),
        note: `context ${percent}% (below ${triggerPercent}% trigger)`,
      };
    }

    const alreadyFresh =
      existing !== undefined && messages.length - existing.covered < RESUMMARISE_EVERY;
    if (existing?.busy === true || alreadyFresh) {
      return {
        block: existing === undefined ? "" : render(existing.summary, percent),
        note: `context ${percent}% (summary current)`,
      };
    }

    const older = messages.slice(0, Math.max(1, messages.length - KEEP_RECENT));
    const state: CompactionState = existing ?? { summary: "", covered: 0, busy: false };
    state.busy = true;
    states.set(key, state);
    try {
      state.summary = await summarise(model, older, state.summary);
      state.covered = messages.length;
    } finally {
      state.busy = false;
    }
    prune();
    return { block: render(state.summary, percent), note: `context ${percent}% -- compacted` };
  } catch {
    return { block: "", note: "" };
  }
}

function render(summary: string, percent: number): string {
  if (summary.trim() === "") return "";
  return [
    OPEN_TAG,
    `The conversation is ${percent}% of the way through the context window, so the earliest turns`,
    "may already have been dropped. This is what happened in them -- treat it as established fact",
    "and do not redo work described here:",
    "",
    summary.trim(),
    CLOSE_TAG,
  ].join("\n");
}

/**
 * Folds the older turns into a compact summary, extending the previous one so
 * that facts from long-dropped turns keep surviving each round.
 */
async function summarise(
  model: MeasurableModel,
  older: ChatMessage[],
  previousSummary: string,
): Promise<string> {
  const body = clamp(
    older.map((m) => `${readRole(m)}: ${readText(m)}`).join("\n"),
    MAX_SOURCE_CHARS,
    "transcript",
  );

  const instruction = [
    "You are compacting a coding conversation so it survives the context window.",
    "Write a dense factual summary that lets the assistant continue without the original turns.",
    "",
    "Include, in this order:",
    "1. What the user is trying to build or fix, in their own terms.",
    "2. Decisions already made and why -- especially ones that must not be revisited.",
    "3. Files created or changed, with paths, and what each change did.",
    "4. Commands run and what they reported (test results, build failures).",
    "5. What is still outstanding.",
    "",
    "Rules: facts only, no praise, no restating these instructions. Keep file paths and error",
    "messages exact. If something is unfinished, say so plainly. Under 400 words.",
  ].join("\n");

  const chat = Chat.from([
    { role: "system", content: instruction },
    {
      role: "user",
      content:
        previousSummary.trim() === ""
          ? `Summarise this conversation:\n\n${body}`
          : `Here is the summary so far:\n\n${previousSummary}\n\nExtend it to cover this fuller transcript, keeping everything above that still matters:\n\n${body}`,
    },
  ]);

  const result = (await model.respond(chat, {
    maxTokens: SUMMARY_TOKENS,
    temperature: 0.2,
  })) as string | { content?: unknown };
  const text =
    typeof result === "string"
      ? result
      : typeof result?.content === "string"
        ? result.content
        : "";
  return clamp(stripReasoning(text), MAX_SUMMARY_CHARS, "summary");
}

function prune(): void {
  while (states.size > MAX_TRACKED_CHATS) {
    const oldest = states.keys().next();
    if (oldest.done === true) break;
    states.delete(oldest.value);
  }
}

import { Chat, type ChatMessage, type PromptPreprocessorController } from "@lmstudio/sdk";
import { clamp, stripReasoning } from "./workspace";

/**
 * Automatic context compaction.
 *
 * A plugin cannot delete earlier turns -- a prompt preprocessor may only rewrite
 * the newest user message. What it CAN do is notice that the conversation is
 * approaching the model's context limit and attach a running summary of the
 * older turns to the newest message. When the host then drops the oldest
 * messages to make room, the summary survives, because it now lives at the end
 * of the conversation rather than the beginning.
 *
 * The net effect matches auto-compaction: work done thirty turns ago is still
 * known to the model after the raw turns have fallen out of the window.
 */

const OPEN_TAG = "<earlier-conversation-summary>";
const CLOSE_TAG = "</earlier-conversation-summary>";

/** Turns kept verbatim at the end; only what precedes them is summarised. */
const KEEP_RECENT = 6;
/** Re-summarise only after the conversation has grown by this many messages. */
const RESUMMARISE_EVERY = 8;
const SUMMARY_TOKENS = 600;
const MAX_SUMMARY_CHARS = 4000;
const MAX_SOURCE_CHARS = 60000;
const MAX_TRACKED_CHATS = 12;

/** The subset of an LLM handle compaction needs, feature-detected at runtime. */
interface MeasurableModel {
  getContextLength?: () => Promise<number>;
  countTokens?: (input: string) => Promise<number>;
  respond: (chat: Chat, opts: { maxTokens: number; temperature: number }) => unknown;
}

interface CompactionState {
  summary: string;
  /** Number of messages already folded into the summary. */
  covered: number;
  /** Guards against two overlapping summarisations in one conversation. */
  busy: boolean;
}

const states = new Map<string, CompactionState>();

/**
 * Conversations carry no id, so identify one by its opening exchange. Stable for
 * the life of a chat, and different chats effectively never collide.
 */
function fingerprint(messages: ChatMessage[]): string {
  // Three messages rather than two: chats often open identically ("make a
  // platformer"), and a collision would leak one chat's summary into another
  // until the next refresh.
  const opening = messages.slice(0, 3).map(readText);
  const shape = opening.map((text) => text.length).join(",");
  return `${shape}|${opening.join("~").slice(0, 300)}`;
}

function readText(message: unknown): string {
  if (typeof message === "string") return message;
  if (message === null || typeof message !== "object") return "";
  const candidate = message as { getText?: () => unknown; content?: unknown; text?: unknown };
  try {
    if (typeof candidate.getText === "function") {
      const value = candidate.getText();
      if (typeof value === "string") return value;
    }
  } catch {
    // Fall through.
  }
  if (typeof candidate.text === "string") return candidate.text;
  if (typeof candidate.content === "string") return candidate.content;
  return "";
}

function readRole(message: unknown): string {
  const candidate = message as { getRole?: () => unknown; role?: unknown };
  try {
    if (typeof candidate?.getRole === "function") {
      const value = candidate.getRole();
      if (typeof value === "string") return value;
    }
  } catch {
    // Fall through.
  }
  return typeof candidate?.role === "string" ? candidate.role : "unknown";
}

export interface CompactionResult {
  /** Block to inject, or "" when there is nothing to add. */
  block: string;
  /** Diagnostic line for the plugin log. */
  note: string;
}

/**
 * Returns a summary block when the conversation is close to filling the context
 * window. Never throws: on any failure the caller simply gets no block.
 */
export async function maybeCompact(
  ctl: PromptPreprocessorController,
  triggerPercent: number,
): Promise<CompactionResult> {
  try {
    const history = await ctl.pullHistory();
    const messages = history.getMessagesArray();
    if (messages.length < KEEP_RECENT + 2) return { block: "", note: "" };

    const key = fingerprint(messages);
    const existing = states.get(key);

    // tokenSource() may hand back a generator handle, which cannot measure
    // tokens. Feature-detect rather than assume a full LLM.
    const model = (await ctl.tokenSource()) as unknown as MeasurableModel;
    if (typeof model.getContextLength !== "function" || typeof model.countTokens !== "function") {
      return { block: "", note: "" };
    }
    const contextLength = await model.getContextLength();
    if (!Number.isFinite(contextLength) || contextLength <= 0) return { block: "", note: "" };

    const transcript = messages.map((m) => `${readRole(m)}: ${readText(m)}`).join("\n");
    const used = await model.countTokens(transcript);
    const percent = Math.round((used / contextLength) * 100);

    // Below the trigger: keep serving any summary already built, so the model
    // does not lose it again the moment usage dips.
    if (percent < triggerPercent) {
      return {
        block: existing === undefined ? "" : render(existing.summary, percent),
        note: `context ${percent}% (below ${triggerPercent}% trigger)`,
      };
    }

    const alreadyFresh =
      existing !== undefined && messages.length - existing.covered < RESUMMARISE_EVERY;
    if (existing?.busy === true || alreadyFresh) {
      return {
        block: existing === undefined ? "" : render(existing.summary, percent),
        note: `context ${percent}% (summary current)`,
      };
    }

    const older = messages.slice(0, Math.max(1, messages.length - KEEP_RECENT));
    const state: CompactionState = existing ?? { summary: "", covered: 0, busy: false };
    state.busy = true;
    states.set(key, state);
    try {
      state.summary = await summarise(model, older, state.summary);
      state.covered = messages.length;
    } finally {
      state.busy = false;
    }
    prune();
    return { block: render(state.summary, percent), note: `context ${percent}% -- compacted` };
  } catch {
    return { block: "", note: "" };
  }
}

function render(summary: string, percent: number): string {
  if (summary.trim() === "") return "";
  return [
    OPEN_TAG,
    `The conversation is ${percent}% of the way through the context window, so the earliest turns`,
    "may already have been dropped. This is what happened in them -- treat it as established fact",
    "and do not redo work described here:",
    "",
    summary.trim(),
    CLOSE_TAG,
  ].join("\n");
}

/**
 * Folds the older turns into a compact summary, extending the previous one so
 * that facts from long-dropped turns keep surviving each round.
 */
async function summarise(
  model: MeasurableModel,
  older: ChatMessage[],
  previousSummary: string,
): Promise<string> {
  const body = clamp(
    older.map((m) => `${readRole(m)}: ${readText(m)}`).join("\n"),
    MAX_SOURCE_CHARS,
    "transcript",
  );

  const instruction = [
    "You are compacting a coding conversation so it survives the context window.",
    "Write a dense factual summary that lets the assistant continue without the original turns.",
    "",
    "Include, in this order:",
    "1. What the user is trying to build or fix, in their own terms.",
    "2. Decisions already made and why -- especially ones that must not be revisited.",
    "3. Files created or changed, with paths, and what each change did.",
    "4. Commands run and what they reported (test results, build failures).",
    "5. What is still outstanding.",
    "",
    "Rules: facts only, no praise, no restating these instructions. Keep file paths and error",
    "messages exact. If something is unfinished, say so plainly. Under 400 words.",
  ].join("\n");

  const chat = Chat.from([
    { role: "system", content: instruction },
    {
      role: "user",
      content:
        previousSummary.trim() === ""
          ? `Summarise this conversation:\n\n${body}`
          : `Here is the summary so far:\n\n${previousSummary}\n\nExtend it to cover this fuller transcript, keeping everything above that still matters:\n\n${body}`,
    },
  ]);

  const result = (await model.respond(chat, {
    maxTokens: SUMMARY_TOKENS,
    temperature: 0.2,
  })) as string | { content?: unknown };
  const text =
    typeof result === "string"
      ? result
      : typeof result?.content === "string"
        ? result.content
        : "";
  return clamp(stripReasoning(text), MAX_SUMMARY_CHARS, "summary");
}

function prune(): void {
  while (states.size > MAX_TRACKED_CHATS) {
    const oldest = states.keys().next();
    if (oldest.done === true) break;
    states.delete(oldest.value);
  }
}