// Runs on every user message.
//
// In a chat with no project it returns the message untouched, unless the chat has just had two
// research-heavy turns in a row (then it starts a project) or the message asks about the
// user's projects (then it attaches the list). In a project chat it applies any hand edit of
// ledger.md, checks sources against the pages read in this chat, and appends the project state:
// the full state when the last full view has left the part of the context kept word for word,
// a short header otherwise, and nothing after the first turn in chats with attached documents.
// It also handles /tl commands.

import type { Chat, ChatMessage, PromptPreprocessorController } from "@lmstudio/sdk";
import * as path from "path";
import { ChatState, chatIdOf, loadChatState, newChatState, saveChatState } from "./chat";
import { applyHandEdit, healthIssues, reminderLine, staleness, verifySources } from "./checks";
import { Exchange, handleCommand, isCommand } from "./commands";
import { configSchematics } from "./config";
import { asksAboutProjects, BLOCK_CLOSE, describeChanges, FULL_OPEN, renderFull, renderHeader, renderProjectList } from "./digest";
import { ownText, turnsOf, workCalls } from "./history";
import { hasIndex, historyTokens, updateIndex, withoutReasoning } from "./recall";
import { createLedger, currentStep, Ledger, loadLedger, loadMeta, progress, readEvents, saveMeta } from "./ledger";
import { appendLine, rootDir } from "./store";

interface Settings {
  recall: boolean;
  recallThresholdTokens: number;
  autoStart: boolean;
  heavyTurnCalls: number;
  staleAfterTurns: number;
  fullDigestTokens: number;
  verbatimWindowTokens: number;
  showStatus: boolean;
}

const DEFAULTS: Settings = { recall: true, recallThresholdTokens: 24000, autoStart: true, heavyTurnCalls: 3, staleAfterTurns: 2, fullDigestTokens: 600, verbatimWindowTokens: 6000, showStatus: true };

function settings(ctl: PromptPreprocessorController): Settings {
  try {
    const c = ctl.getPluginConfig(configSchematics);
    const get = <K extends keyof Settings>(k: K): Settings[K] => {
      const v = c.get(k) as Settings[K] | undefined;
      return v === undefined || v === null ? DEFAULTS[k] : v;
    };
    return {
      recall: get("recall"),
      recallThresholdTokens: get("recallThresholdTokens"),
      autoStart: get("autoStart"),
      heavyTurnCalls: get("heavyTurnCalls"),
      staleAfterTurns: get("staleAfterTurns"),
      fullDigestTokens: get("fullDigestTokens"),
      verbatimWindowTokens: get("verbatimWindowTokens"),
      showStatus: get("showStatus"),
    };
  } catch {
    return DEFAULTS;
  }
}

function showStatus(ctl: PromptPreprocessorController, on: boolean, text: string): void {
  if (!on) return;
  try {
    ctl.createStatus({ status: "done", text });
  } catch {
    // A status line is a convenience; never let it break the message.
  }
}

function messageChars(m: ChatMessage): number {
  return m.getText().length + m.getToolCallResults().reduce((n, r) => n + r.content.length, 0);
}

// Tokens of history after the last message that carried the full state, or Infinity if none did.
export function tokensSinceLastFull(history: Chat): number {
  const msgs = history.getMessagesArray();
  let chars = 0;
  for (let i = msgs.length - 1; i >= 0; i--) {
    const m = msgs[i];
    if (m.getRole() === "user" && m.getText().includes(FULL_OPEN)) return Math.ceil(chars / 4);
    chars += messageChars(m);
  }
  return Infinity;
}

function statusText(l: Ledger, kind: string): string {
  const p = progress(l);
  const cur = currentStep(l);
  const where = !p.total ? "no plan yet" : cur ? `step ${cur.id} of ${p.total}` : `${p.done} of ${p.total} done`;
  return `throughLine: ${l.name} | ${where}${kind ? ` | ${kind}` : ""}`;
}

// True when the chat has a document attached (not just images). rag-v1 searches documents
// using the full text of each user message, so throughLine keeps its text out of those chats.
function hasDocuments(ctl: PromptPreprocessorController, history: Chat, userMessage: ChatMessage): boolean {
  if (!history.hasFiles() && !userMessage.hasFiles()) return false;
  try {
    const files = [...history.getAllFiles(ctl.client), ...userMessage.getFiles(ctl.client)];
    return files.some((f) => f.type !== "image");
  } catch {
    return true;
  }
}

// Records whether rag-v1 had already rewritten this message when throughLine received it.
// One boolean per document turn; never any text.
function recordOrder(userMessage: ChatMessage): void {
  const t = userMessage.getText();
  const ragRanFirst = t.includes("User Query:") || t.includes("Enriched Context Generation");
  try {
    appendLine(path.join(rootDir(), "diagnostics.jsonl"), JSON.stringify({ ts: new Date().toISOString(), event: "preprocessor_order", ragRanFirst }));
  } catch {
    // Diagnostics are optional.
  }
}

export const DOCUMENT_CHAT_NOTE = "This chat has attached documents, so this state is not repeated on later messages. Use ledger_view when you need it.";
export const AUTO_START_NOTE = "This project started automatically because the last two turns were research-heavy. If the user did not want it, they can send /tl off.";

export async function preprocess(ctl: PromptPreprocessorController, userMessage: ChatMessage): Promise<string | ChatMessage> {
  const chatId = chatIdOf(ctl);
  if (!chatId) return userMessage;
  const text = userMessage.getText();
  const cfg = settings(ctl);

  if (isCommand(text)) {
    const history = await ctl.pullHistory();
    const msgs = history.getMessagesArray();
    const userTexts = msgs.filter((m) => m.getRole() === "user").map((m) => m.getText());
    const recent: Exchange[] = msgs
      .filter((m) => m.getRole() === "user" || m.getRole() === "assistant")
      .map((m) => ({ role: m.getRole() as "user" | "assistant", text: m.getRole() === "user" ? ownText(m.getText()) : withoutReasoning(m.getText()) }))
      .filter((e) => e.text && !e.text.startsWith("[throughLine"))
      .slice(-8);
    const result = handleCommand(chatId, text, userTexts, cfg.fullDigestTokens, recent);
    showStatus(ctl, true, result.status);
    const out = userMessage.asMutableCopy();
    out.replaceText(result.forModel);
    return out;
  }

  let state: ChatState | null = loadChatState(chatId);
  let ledger = state?.project ? loadLedger(state.project) : null;
  if (state?.project && !ledger) {
    state.project = undefined;
    saveChatState(state);
  }
  const asks = asksAboutProjects(text);
  const mayAutoStart = cfg.autoStart && !ledger && !state?.autostartOff;
  const mayRecall = cfg.recall && !state?.recallOff;
  if (!ledger && !mayAutoStart && !asks && !mayRecall) return userMessage;

  const history = await ctl.pullHistory();
  const turns = turnsOf(history);

  // Recall: index a long chat, then keep the index current. Nothing is copied before the threshold.
  let recallStarted = false;
  if (mayRecall) {
    try {
      const indexed = hasIndex(chatId);
      if (indexed || historyTokens(history) >= cfg.recallThresholdTokens) {
        updateIndex(chatId, history, cfg.verbatimWindowTokens);
        recallStarted = !indexed;
      }
    } catch {
      // Recall is an aid; a failure here must not affect the message.
    }
  }
  if (recallStarted) showStatus(ctl, cfg.showStatus, "throughLine: this chat is long, so recall is on (/tl recall off to stop)");

  // Automatic start: the last two complete turns each used several tools for real work.
  let started = false;
  if (!ledger && mayAutoStart) {
    const lastTwo = turns.slice(-2);
    if (lastTwo.length === 2 && lastTwo.every((t) => workCalls(t) >= cfg.heavyTurnCalls)) {
      const asks = lastTwo.map((t) => t.userText).filter((g) => g && !g.startsWith("/") && !g.startsWith("[throughLine"));
      const goal = [...new Set(asks)].join(" ");
      ledger = createLedger(goal.slice(0, 600).trim(), goal ? "user_message" : "none", chatId, { autoStarted: true });
      state = state ?? newChatState(chatId);
      state.project = ledger.name;
      state.lastSeq = ledger.seq;
      saveChatState(state);
      started = true;
    }
  }

  const out = userMessage.asMutableCopy();
  if (!ledger) {
    if (!asks) return userMessage;
    out.appendText(`\n\n${renderProjectList()}`);
    showStatus(ctl, cfg.showStatus, "throughLine: project list attached");
    return out;
  }
  const st = state!;

  // Bookkeeping for this turn, from what happened since this chat last saw the ledger.
  const meta = loadMeta(ledger.name);
  meta.turns++;
  saveMeta(ledger.name, meta);
  st.turn = (st.turn ?? 0) + 1;
  st.turnAt = new Date().toISOString();
  const sinceLast = readEvents(ledger.name, st.lastSeq);
  const modelChanged = sinceLast.some((e) => e.actor === "model");
  const edit = applyHandEdit(ledger);
  verifySources(ledger, chatId, turns);
  const stale = staleness(st, st.turn, turns.at(-1), modelChanged, cfg.staleAfterTurns);
  st.staleTurns = stale.staleTurns;
  if (stale.remind) st.lastReminderTurn = st.turn;

  const extras: string[] = [];
  if (edit.kind === "applied") extras.push(`Notice: the user edited the ledger by hand (${edit.summary}). This state includes the edit.`);
  if (edit.kind === "failed") extras.push(`Notice: the user's edit to ledger.md could not be read (${edit.reason}). It was saved as ${edit.savedAs} and the ledger is unchanged; tell the user.`);
  const issues = healthIssues(ledger, meta, sinceLast);
  if (issues.length) extras.push(`Health: ${issues.join("; ")}.`);
  if (stale.remind) extras.push(reminderLine(stale.staleTurns));
  if (started) extras.push(AUTO_START_NOTE);

  const sinceFull = tokensSinceLastFull(history);
  const documents = hasDocuments(ctl, history, userMessage);
  const mustShowFull = started || edit.kind !== "none" || sinceFull === Infinity;

  let block: string | null;
  let kind = "";
  if (documents) {
    recordOrder(userMessage);
    if (mustShowFull) {
      block = renderFull(ledger, cfg.fullDigestTokens, [...extras, DOCUMENT_CHAT_NOTE]);
      kind = "full state shown";
    } else {
      // The project's state was already shown in this chat; stay out of rag-v1's search.
      block = null;
      kind = "quiet: documents attached";
    }
  } else if (mustShowFull || sinceFull > cfg.verbatimWindowTokens) {
    block = renderFull(ledger, cfg.fullDigestTokens, extras);
    kind = "full state shown";
  } else {
    block = renderHeader(ledger, describeChanges(sinceLast), extras);
  }
  if (stale.remind && block) kind = kind ? `${kind} | reminder sent` : "reminder sent";
  if (edit.kind === "applied") kind = kind ? `${kind} | your edit applied` : "your edit applied";
  if (edit.kind === "failed") kind = kind ? `${kind} | your edit could not be read` : "your edit could not be read";

  st.lastSeq = ledger.seq;
  saveChatState(st);
  showStatus(ctl, cfg.showStatus, started ? `throughLine: started ${ledger.name} automatically | /tl off to stop` : statusText(ledger, kind));

  const parts = [block, asks ? renderProjectList(ledger.name) : null].filter((p): p is string => !!p);
  if (!parts.length) return userMessage;
  out.appendText(`\n\n${parts.join("\n\n")}`);
  return out;
}

export { BLOCK_CLOSE };

// Runs on every user message.
//
// In a chat with no project it returns the message untouched, unless the chat has just had two
// research-heavy turns in a row (then it starts a project) or the message asks about the
// user's projects (then it attaches the list). In a project chat it applies any hand edit of
// ledger.md, checks sources against the pages read in this chat, and appends the project state:
// the full state when the last full view has left the part of the context kept word for word,
// a short header otherwise, and nothing after the first turn in chats with attached documents.
// It also handles /tl commands.

import type { Chat, ChatMessage, PromptPreprocessorController } from "@lmstudio/sdk";
import * as path from "path";
import { ChatState, chatIdOf, loadChatState, newChatState, saveChatState } from "./chat";
import { applyHandEdit, healthIssues, reminderLine, staleness, verifySources } from "./checks";
import { Exchange, handleCommand, isCommand } from "./commands";
import { configSchematics } from "./config";
import { asksAboutProjects, BLOCK_CLOSE, describeChanges, FULL_OPEN, renderFull, renderHeader, renderProjectList } from "./digest";
import { ownText, turnsOf, workCalls } from "./history";
import { hasIndex, historyTokens, updateIndex, withoutReasoning } from "./recall";
import { createLedger, currentStep, Ledger, loadLedger, loadMeta, progress, readEvents, saveMeta } from "./ledger";
import { appendLine, rootDir } from "./store";

interface Settings {
  recall: boolean;
  recallThresholdTokens: number;
  autoStart: boolean;
  heavyTurnCalls: number;
  staleAfterTurns: number;
  fullDigestTokens: number;
  verbatimWindowTokens: number;
  showStatus: boolean;
}

const DEFAULTS: Settings = { recall: true, recallThresholdTokens: 24000, autoStart: true, heavyTurnCalls: 3, staleAfterTurns: 2, fullDigestTokens: 600, verbatimWindowTokens: 6000, showStatus: true };

function settings(ctl: PromptPreprocessorController): Settings {
  try {
    const c = ctl.getPluginConfig(configSchematics);
    const get = <K extends keyof Settings>(k: K): Settings[K] => {
      const v = c.get(k) as Settings[K] | undefined;
      return v === undefined || v === null ? DEFAULTS[k] : v;
    };
    return {
      recall: get("recall"),
      recallThresholdTokens: get("recallThresholdTokens"),
      autoStart: get("autoStart"),
      heavyTurnCalls: get("heavyTurnCalls"),
      staleAfterTurns: get("staleAfterTurns"),
      fullDigestTokens: get("fullDigestTokens"),
      verbatimWindowTokens: get("verbatimWindowTokens"),
      showStatus: get("showStatus"),
    };
  } catch {
    return DEFAULTS;
  }
}

function showStatus(ctl: PromptPreprocessorController, on: boolean, text: string): void {
  if (!on) return;
  try {
    ctl.createStatus({ status: "done", text });
  } catch {
    // A status line is a convenience; never let it break the message.
  }
}

function messageChars(m: ChatMessage): number {
  return m.getText().length + m.getToolCallResults().reduce((n, r) => n + r.content.length, 0);
}

// Tokens of history after the last message that carried the full state, or Infinity if none did.
export function tokensSinceLastFull(history: Chat): number {
  const msgs = history.getMessagesArray();
  let chars = 0;
  for (let i = msgs.length - 1; i >= 0; i--) {
    const m = msgs[i];
    if (m.getRole() === "user" && m.getText().includes(FULL_OPEN)) return Math.ceil(chars / 4);
    chars += messageChars(m);
  }
  return Infinity;
}

function statusText(l: Ledger, kind: string): string {
  const p = progress(l);
  const cur = currentStep(l);
  const where = !p.total ? "no plan yet" : cur ? `step ${cur.id} of ${p.total}` : `${p.done} of ${p.total} done`;
  return `throughLine: ${l.name} | ${where}${kind ? ` | ${kind}` : ""}`;
}

// True when the chat has a document attached (not just images). rag-v1 searches documents
// using the full text of each user message, so throughLine keeps its text out of those chats.
function hasDocuments(ctl: PromptPreprocessorController, history: Chat, userMessage: ChatMessage): boolean {
  if (!history.hasFiles() && !userMessage.hasFiles()) return false;
  try {
    const files = [...history.getAllFiles(ctl.client), ...userMessage.getFiles(ctl.client)];
    return files.some((f) => f.type !== "image");
  } catch {
    return true;
  }
}

// Records whether rag-v1 had already rewritten this message when throughLine received it.
// One boolean per document turn; never any text.
function recordOrder(userMessage: ChatMessage): void {
  const t = userMessage.getText();
  const ragRanFirst = t.includes("User Query:") || t.includes("Enriched Context Generation");
  try {
    appendLine(path.join(rootDir(), "diagnostics.jsonl"), JSON.stringify({ ts: new Date().toISOString(), event: "preprocessor_order", ragRanFirst }));
  } catch {
    // Diagnostics are optional.
  }
}

export const DOCUMENT_CHAT_NOTE = "This chat has attached documents, so this state is not repeated on later messages. Use ledger_view when you need it.";
export const AUTO_START_NOTE = "This project started automatically because the last two turns were research-heavy. If the user did not want it, they can send /tl off.";

export async function preprocess(ctl: PromptPreprocessorController, userMessage: ChatMessage): Promise<string | ChatMessage> {
  const chatId = chatIdOf(ctl);
  if (!chatId) return userMessage;
  const text = userMessage.getText();
  const cfg = settings(ctl);

  if (isCommand(text)) {
    const history = await ctl.pullHistory();
    const msgs = history.getMessagesArray();
    const userTexts = msgs.filter((m) => m.getRole() === "user").map((m) => m.getText());
    const recent: Exchange[] = msgs
      .filter((m) => m.getRole() === "user" || m.getRole() === "assistant")
      .map((m) => ({ role: m.getRole() as "user" | "assistant", text: m.getRole() === "user" ? ownText(m.getText()) : withoutReasoning(m.getText()) }))
      .filter((e) => e.text && !e.text.startsWith("[throughLine"))
      .slice(-8);
    const result = handleCommand(chatId, text, userTexts, cfg.fullDigestTokens, recent);
    showStatus(ctl, true, result.status);
    const out = userMessage.asMutableCopy();
    out.replaceText(result.forModel);
    return out;
  }

  let state: ChatState | null = loadChatState(chatId);
  let ledger = state?.project ? loadLedger(state.project) : null;
  if (state?.project && !ledger) {
    state.project = undefined;
    saveChatState(state);
  }
  const asks = asksAboutProjects(text);
  const mayAutoStart = cfg.autoStart && !ledger && !state?.autostartOff;
  const mayRecall = cfg.recall && !state?.recallOff;
  if (!ledger && !mayAutoStart && !asks && !mayRecall) return userMessage;

  const history = await ctl.pullHistory();
  const turns = turnsOf(history);

  // Recall: index a long chat, then keep the index current. Nothing is copied before the threshold.
  let recallStarted = false;
  if (mayRecall) {
    try {
      const indexed = hasIndex(chatId);
      if (indexed || historyTokens(history) >= cfg.recallThresholdTokens) {
        updateIndex(chatId, history, cfg.verbatimWindowTokens);
        recallStarted = !indexed;
      }
    } catch {
      // Recall is an aid; a failure here must not affect the message.
    }
  }
  if (recallStarted) showStatus(ctl, cfg.showStatus, "throughLine: this chat is long, so recall is on (/tl recall off to stop)");

  // Automatic start: the last two complete turns each used several tools for real work.
  let started = false;
  if (!ledger && mayAutoStart) {
    const lastTwo = turns.slice(-2);
    if (lastTwo.length === 2 && lastTwo.every((t) => workCalls(t) >= cfg.heavyTurnCalls)) {
      const asks = lastTwo.map((t) => t.userText).filter((g) => g && !g.startsWith("/") && !g.startsWith("[throughLine"));
      const goal = [...new Set(asks)].join(" ");
      ledger = createLedger(goal.slice(0, 600).trim(), goal ? "user_message" : "none", chatId, { autoStarted: true });
      state = state ?? newChatState(chatId);
      state.project = ledger.name;
      state.lastSeq = ledger.seq;
      saveChatState(state);
      started = true;
    }
  }

  const out = userMessage.asMutableCopy();
  if (!ledger) {
    if (!asks) return userMessage;
    out.appendText(`\n\n${renderProjectList()}`);
    showStatus(ctl, cfg.showStatus, "throughLine: project list attached");
    return out;
  }
  const st = state!;

  // Bookkeeping for this turn, from what happened since this chat last saw the ledger.
  const meta = loadMeta(ledger.name);
  meta.turns++;
  saveMeta(ledger.name, meta);
  st.turn = (st.turn ?? 0) + 1;
  st.turnAt = new Date().toISOString();
  const sinceLast = readEvents(ledger.name, st.lastSeq);
  const modelChanged = sinceLast.some((e) => e.actor === "model");
  const edit = applyHandEdit(ledger);
  verifySources(ledger, chatId, turns);
  const stale = staleness(st, st.turn, turns.at(-1), modelChanged, cfg.staleAfterTurns);
  st.staleTurns = stale.staleTurns;
  if (stale.remind) st.lastReminderTurn = st.turn;

  const extras: string[] = [];
  if (edit.kind === "applied") extras.push(`Notice: the user edited the ledger by hand (${edit.summary}). This state includes the edit.`);
  if (edit.kind === "failed") extras.push(`Notice: the user's edit to ledger.md could not be read (${edit.reason}). It was saved as ${edit.savedAs} and the ledger is unchanged; tell the user.`);
  const issues = healthIssues(ledger, meta, sinceLast);
  if (issues.length) extras.push(`Health: ${issues.join("; ")}.`);
  if (stale.remind) extras.push(reminderLine(stale.staleTurns));
  if (started) extras.push(AUTO_START_NOTE);

  const sinceFull = tokensSinceLastFull(history);
  const documents = hasDocuments(ctl, history, userMessage);
  const mustShowFull = started || edit.kind !== "none" || sinceFull === Infinity;

  let block: string | null;
  let kind = "";
  if (documents) {
    recordOrder(userMessage);
    if (mustShowFull) {
      block = renderFull(ledger, cfg.fullDigestTokens, [...extras, DOCUMENT_CHAT_NOTE]);
      kind = "full state shown";
    } else {
      // The project's state was already shown in this chat; stay out of rag-v1's search.
      block = null;
      kind = "quiet: documents attached";
    }
  } else if (mustShowFull || sinceFull > cfg.verbatimWindowTokens) {
    block = renderFull(ledger, cfg.fullDigestTokens, extras);
    kind = "full state shown";
  } else {
    block = renderHeader(ledger, describeChanges(sinceLast), extras);
  }
  if (stale.remind && block) kind = kind ? `${kind} | reminder sent` : "reminder sent";
  if (edit.kind === "applied") kind = kind ? `${kind} | your edit applied` : "your edit applied";
  if (edit.kind === "failed") kind = kind ? `${kind} | your edit could not be read` : "your edit could not be read";

  st.lastSeq = ledger.seq;
  saveChatState(st);
  showStatus(ctl, cfg.showStatus, started ? `throughLine: started ${ledger.name} automatically | /tl off to stop` : statusText(ledger, kind));

  const parts = [block, asks ? renderProjectList(ledger.name) : null].filter((p): p is string => !!p);
  if (!parts.length) return userMessage;
  out.appendText(`\n\n${parts.join("\n\n")}`);
  return out;
}

export { BLOCK_CLOSE };