// Recall: a searchable copy of a long chat's stored history, so the model can find the original
// wording of anything compression has summarized away.
//
// The index is built only once a chat's history passes the recall threshold, kept in the chat's
// own folder, and extended by the new messages on each turn. If earlier messages change (an
// edit, a regenerated reply, a deleted message), it is rebuilt from the full history.

import type { Chat } from "@lmstudio/sdk";
import * as fs from "fs";
import * as path from "path";
import { hash } from "./ledger";
import { ownText } from "./history";
import { appendLine, chatDir, readJson, writeJson } from "./store";
import { clip } from "./text";

export interface Passage {
  m: number; // message index in the history
  turn: number; // user turn the message belongs to (1-based)
  role: string; // user, assistant, or tool:<name>
  ts: string; // when throughLine first indexed it
  text: string;
}

interface IndexState {
  hashes: string[]; // a short fingerprint of each indexed message, to notice edits anywhere
  skipFrom: number; // first message still in the model's view word for word
  updated: string;
}

const fingerprint = (m: { getText(): string; getToolCallResults(): Array<{ content: string }> }) =>
  hash(m.getText() + "\u0000" + m.getToolCallResults().map((r) => r.content).join("\u0000")).slice(0, 12);

const PASSAGE_CHARS = 900;
export const RESULT_CHARS = 6000;

const dir = (chatId: string) => path.join(chatDir(chatId), "recall");
const passagesFile = (chatId: string) => path.join(dir(chatId), "passages.jsonl");
const stateFile = (chatId: string) => path.join(dir(chatId), "state.json");

export function historyTokens(history: Chat): number {
  let chars = 0;
  for (const m of history.getMessagesArray()) {
    chars += m.getText().length;
    for (const r of m.getToolCallResults()) chars += r.content.length;
  }
  return Math.ceil(chars / 4);
}

export function hasIndex(chatId: string): boolean {
  return fs.existsSync(stateFile(chatId));
}

export function deleteIndex(chatId: string): void {
  fs.rmSync(dir(chatId), { recursive: true, force: true });
}

// Splits text into passages of roughly PASSAGE_CHARS, on paragraph and then sentence boundaries.
export function splitPassages(text: string): string[] {
  const out: string[] = [];
  let cur = "";
  const push = () => { if (cur.trim()) out.push(cur.trim()); cur = ""; };
  for (const para of text.split(/\n\s*\n/)) {
    const pieces = para.length > PASSAGE_CHARS ? para.match(/[^.!?\n]+[.!?]*\s*|\n/g) ?? [para] : [para];
    for (const piece of pieces) {
      if (cur.length + piece.length > PASSAGE_CHARS) push();
      if (piece.length > PASSAGE_CHARS) {
        for (let i = 0; i < piece.length; i += PASSAGE_CHARS) out.push(piece.slice(i, i + PASSAGE_CHARS).trim());
        continue;
      }
      cur += (cur && !cur.endsWith("\n") ? "\n\n" : "") + piece;
    }
  }
  push();
  return out.filter(Boolean);
}

// Reasoning the model wrote before its answer, where LM Studio stores it in the text.
export function withoutReasoning(text: string): string {
  return text.replace(/<think>[\s\S]*?<\/think>/g, "").replace(/^[\s\S]*LM_STUDIO_INTERNAL_LSEP_SYNTHETIC_REASONING_END_[0-9a-f]+_*/, "").trim();
}

function messagePassages(history: Chat, from: number, turnAtFrom: number): Passage[] {
  const msgs = history.getMessagesArray();
  const ts = new Date().toISOString();
  const out: Passage[] = [];
  let turn = turnAtFrom;
  const names = new Map<string, string>();
  for (const m of msgs) for (const r of m.getToolCallRequests()) if (r.id) names.set(r.id, r.name);
  for (let i = from; i < msgs.length; i++) {
    const m = msgs[i];
    const role = m.getRole();
    if (role === "system") continue;
    if (role === "user") turn++;
    const t = Math.max(turn, 1);
    if (role === "user") {
      for (const p of splitPassages(ownText(m.getText()))) out.push({ m: i, turn: t, role: "user", ts, text: p });
    } else if (role === "assistant") {
      for (const p of splitPassages(withoutReasoning(m.getText()))) out.push({ m: i, turn: t, role: "assistant", ts, text: p });
    }
    for (const r of m.getToolCallResults()) {
      const name = (r.toolCallId && names.get(r.toolCallId)) || "tool";
      if (name.startsWith("ledger_")) continue;
      for (const p of splitPassages(r.content)) out.push({ m: i, turn: t, role: `tool:${name}`, ts, text: p });
    }
  }
  return out;
}

function turnsBefore(history: Chat, index: number): number {
  return history.getMessagesArray().slice(0, index).filter((m) => m.getRole() === "user").length;
}

// Brings the index up to date with the stored history: messages after the last indexed one are
// added, and if an earlier message was edited or removed, everything from there is redone.
// Returns the number of passages added.
export function updateIndex(chatId: string, history: Chat, verbatimTokens: number): number {
  const msgs = history.getMessagesArray();
  const prints = msgs.map(fingerprint);
  const st = readJson<IndexState>(stateFile(chatId));
  let from = 0;
  if (st?.hashes) {
    while (from < st.hashes.length && from < prints.length && st.hashes[from] === prints[from]) from++;
    if (from < st.hashes.length) {
      // History changed at message `from`: keep only passages from before it.
      const kept = loadPassages(chatId).filter((p) => p.m < from);
      fs.mkdirSync(dir(chatId), { recursive: true });
      fs.writeFileSync(passagesFile(chatId), kept.map((p) => JSON.stringify(p) + "\n").join(""));
    }
  }
  const added = messagePassages(history, from, turnsBefore(history, from));
  for (const p of added) appendLine(passagesFile(chatId), JSON.stringify(p));
  writeJson(stateFile(chatId), { hashes: prints, skipFrom: verbatimStart(history, verbatimTokens), updated: new Date().toISOString() } satisfies IndexState);
  return added.length;
}

function loadPassages(chatId: string): Passage[] {
  try {
    return fs.readFileSync(passagesFile(chatId), "utf8").split("\n").filter(Boolean).map((l) => JSON.parse(l) as Passage);
  } catch {
    return [];
  }
}

const STOP = new Set("a an and are as at be but by for from has have he her his i in is it its me my of on or our she so that the their them they this to was we were what when where which who why will with you your".split(" "));
// Words, with compounds ("browser_unavailable", "v26.12.0", "read-only") also split into their
// parts, so a search for one part finds the compound.
const tokenize = (t: string): string[] => {
  const out: string[] = [];
  for (const w of t.toLowerCase().normalize("NFKD").replace(/[\u0300-\u036f]/g, "").match(/[a-z0-9][a-z0-9._-]*[a-z0-9]|[a-z0-9]/g) ?? []) {
    if (!STOP.has(w)) out.push(w);
    if (/[._-]/.test(w)) for (const part of w.split(/[._-]+/)) if (part.length > 1 && !STOP.has(part)) out.push(part);
  }
  return out;
};

// BM25 ranking of passages against the query. Passages from the most recent messages, which
// are still in the model's view word for word, are skipped.
export function search(chatId: string, query: string, opts: { limit?: number } = {}): Passage[] {
  const skipFrom = readJson<IndexState>(stateFile(chatId))?.skipFrom ?? Infinity;
  const all = loadPassages(chatId).filter((p) => p.m < skipFrom);
  const q = [...new Set(tokenize(query))];
  if (!q.length || !all.length) return [];
  const docs = all.map((p) => tokenize(p.text));
  const avg = docs.reduce((n, d) => n + d.length, 0) / docs.length || 1;
  const df = new Map<string, number>();
  for (const d of docs) for (const w of new Set(d)) df.set(w, (df.get(w) ?? 0) + 1);
  const k1 = 1.4, b = 0.75, N = docs.length;
  const scored = docs.map((d, i) => {
    const tf = new Map<string, number>();
    for (const w of d) tf.set(w, (tf.get(w) ?? 0) + 1);
    let s = 0;
    for (const w of q) {
      const f = tf.get(w);
      if (!f) continue;
      const idf = Math.log(1 + (N - (df.get(w) ?? 0) + 0.5) / ((df.get(w) ?? 0) + 0.5));
      s += idf * ((f * (k1 + 1)) / (f + k1 * (1 - b + (b * d.length) / avg)));
    }
    return { p: all[i], s };
  });
  return scored.filter((x) => x.s > 0).sort((a, b2) => b2.s - a.s).slice(0, opts.limit ?? 4).map((x) => x.p);
}

export function formatResults(results: Passage[]): string {
  if (!results.length) return "status: no_results\nnext: Nothing earlier in this chat matches. Try other words, or tell the user you could not find it.";
  const lines = ["status: ok"];
  let used = 0;
  for (const p of results.sort((a, b) => a.m - b.m)) {
    const who = p.role === "user" ? "the user" : p.role === "assistant" ? "the assistant" : `tool ${p.role.slice(5)}`;
    const entry = `[turn ${p.turn}, ${who}, ${p.ts.slice(0, 16).replace("T", " ")} UTC]\n${clip(p.text, Math.min(1600, RESULT_CHARS - used))}`;
    if (used + entry.length > RESULT_CHARS) break;
    lines.push(entry);
    used += entry.length;
  }
  lines.push("next: These are the original words from earlier in this chat. Cite the turn when you rely on them, for example as a finding's source (\"turn 12\").");
  return lines.join("\n\n");
}

// The first message index whose content is still within the verbatim window, counting back
// from the newest message. Passages from there on are in the model's view already.
export function verbatimStart(history: Chat, windowTokens: number): number {
  const msgs = history.getMessagesArray();
  let chars = 0;
  for (let i = msgs.length - 1; i >= 0; i--) {
    chars += msgs[i].getText().length + msgs[i].getToolCallResults().reduce((n, r) => n + r.content.length, 0);
    if (chars / 4 > windowTokens) return i + 1;
  }
  return 0;
}

// Recall: a searchable copy of a long chat's stored history, so the model can find the original
// wording of anything compression has summarized away.
//
// The index is built only once a chat's history passes the recall threshold, kept in the chat's
// own folder, and extended by the new messages on each turn. If earlier messages change (an
// edit, a regenerated reply, a deleted message), it is rebuilt from the full history.

import type { Chat } from "@lmstudio/sdk";
import * as fs from "fs";
import * as path from "path";
import { hash } from "./ledger";
import { ownText } from "./history";
import { appendLine, chatDir, readJson, writeJson } from "./store";
import { clip } from "./text";

export interface Passage {
  m: number; // message index in the history
  turn: number; // user turn the message belongs to (1-based)
  role: string; // user, assistant, or tool:<name>
  ts: string; // when throughLine first indexed it
  text: string;
}

interface IndexState {
  hashes: string[]; // a short fingerprint of each indexed message, to notice edits anywhere
  skipFrom: number; // first message still in the model's view word for word
  updated: string;
}

const fingerprint = (m: { getText(): string; getToolCallResults(): Array<{ content: string }> }) =>
  hash(m.getText() + "\u0000" + m.getToolCallResults().map((r) => r.content).join("\u0000")).slice(0, 12);

const PASSAGE_CHARS = 900;
export const RESULT_CHARS = 6000;

const dir = (chatId: string) => path.join(chatDir(chatId), "recall");
const passagesFile = (chatId: string) => path.join(dir(chatId), "passages.jsonl");
const stateFile = (chatId: string) => path.join(dir(chatId), "state.json");

export function historyTokens(history: Chat): number {
  let chars = 0;
  for (const m of history.getMessagesArray()) {
    chars += m.getText().length;
    for (const r of m.getToolCallResults()) chars += r.content.length;
  }
  return Math.ceil(chars / 4);
}

export function hasIndex(chatId: string): boolean {
  return fs.existsSync(stateFile(chatId));
}

export function deleteIndex(chatId: string): void {
  fs.rmSync(dir(chatId), { recursive: true, force: true });
}

// Splits text into passages of roughly PASSAGE_CHARS, on paragraph and then sentence boundaries.
export function splitPassages(text: string): string[] {
  const out: string[] = [];
  let cur = "";
  const push = () => { if (cur.trim()) out.push(cur.trim()); cur = ""; };
  for (const para of text.split(/\n\s*\n/)) {
    const pieces = para.length > PASSAGE_CHARS ? para.match(/[^.!?\n]+[.!?]*\s*|\n/g) ?? [para] : [para];
    for (const piece of pieces) {
      if (cur.length + piece.length > PASSAGE_CHARS) push();
      if (piece.length > PASSAGE_CHARS) {
        for (let i = 0; i < piece.length; i += PASSAGE_CHARS) out.push(piece.slice(i, i + PASSAGE_CHARS).trim());
        continue;
      }
      cur += (cur && !cur.endsWith("\n") ? "\n\n" : "") + piece;
    }
  }
  push();
  return out.filter(Boolean);
}

// Reasoning the model wrote before its answer, where LM Studio stores it in the text.
export function withoutReasoning(text: string): string {
  return text.replace(/<think>[\s\S]*?<\/think>/g, "").replace(/^[\s\S]*LM_STUDIO_INTERNAL_LSEP_SYNTHETIC_REASONING_END_[0-9a-f]+_*/, "").trim();
}

function messagePassages(history: Chat, from: number, turnAtFrom: number): Passage[] {
  const msgs = history.getMessagesArray();
  const ts = new Date().toISOString();
  const out: Passage[] = [];
  let turn = turnAtFrom;
  const names = new Map<string, string>();
  for (const m of msgs) for (const r of m.getToolCallRequests()) if (r.id) names.set(r.id, r.name);
  for (let i = from; i < msgs.length; i++) {
    const m = msgs[i];
    const role = m.getRole();
    if (role === "system") continue;
    if (role === "user") turn++;
    const t = Math.max(turn, 1);
    if (role === "user") {
      for (const p of splitPassages(ownText(m.getText()))) out.push({ m: i, turn: t, role: "user", ts, text: p });
    } else if (role === "assistant") {
      for (const p of splitPassages(withoutReasoning(m.getText()))) out.push({ m: i, turn: t, role: "assistant", ts, text: p });
    }
    for (const r of m.getToolCallResults()) {
      const name = (r.toolCallId && names.get(r.toolCallId)) || "tool";
      if (name.startsWith("ledger_")) continue;
      for (const p of splitPassages(r.content)) out.push({ m: i, turn: t, role: `tool:${name}`, ts, text: p });
    }
  }
  return out;
}

function turnsBefore(history: Chat, index: number): number {
  return history.getMessagesArray().slice(0, index).filter((m) => m.getRole() === "user").length;
}

// Brings the index up to date with the stored history: messages after the last indexed one are
// added, and if an earlier message was edited or removed, everything from there is redone.
// Returns the number of passages added.
export function updateIndex(chatId: string, history: Chat, verbatimTokens: number): number {
  const msgs = history.getMessagesArray();
  const prints = msgs.map(fingerprint);
  const st = readJson<IndexState>(stateFile(chatId));
  let from = 0;
  if (st?.hashes) {
    while (from < st.hashes.length && from < prints.length && st.hashes[from] === prints[from]) from++;
    if (from < st.hashes.length) {
      // History changed at message `from`: keep only passages from before it.
      const kept = loadPassages(chatId).filter((p) => p.m < from);
      fs.mkdirSync(dir(chatId), { recursive: true });
      fs.writeFileSync(passagesFile(chatId), kept.map((p) => JSON.stringify(p) + "\n").join(""));
    }
  }
  const added = messagePassages(history, from, turnsBefore(history, from));
  for (const p of added) appendLine(passagesFile(chatId), JSON.stringify(p));
  writeJson(stateFile(chatId), { hashes: prints, skipFrom: verbatimStart(history, verbatimTokens), updated: new Date().toISOString() } satisfies IndexState);
  return added.length;
}

function loadPassages(chatId: string): Passage[] {
  try {
    return fs.readFileSync(passagesFile(chatId), "utf8").split("\n").filter(Boolean).map((l) => JSON.parse(l) as Passage);
  } catch {
    return [];
  }
}

const STOP = new Set("a an and are as at be but by for from has have he her his i in is it its me my of on or our she so that the their them they this to was we were what when where which who why will with you your".split(" "));
// Words, with compounds ("browser_unavailable", "v26.12.0", "read-only") also split into their
// parts, so a search for one part finds the compound.
const tokenize = (t: string): string[] => {
  const out: string[] = [];
  for (const w of t.toLowerCase().normalize("NFKD").replace(/[\u0300-\u036f]/g, "").match(/[a-z0-9][a-z0-9._-]*[a-z0-9]|[a-z0-9]/g) ?? []) {
    if (!STOP.has(w)) out.push(w);
    if (/[._-]/.test(w)) for (const part of w.split(/[._-]+/)) if (part.length > 1 && !STOP.has(part)) out.push(part);
  }
  return out;
};

// BM25 ranking of passages against the query. Passages from the most recent messages, which
// are still in the model's view word for word, are skipped.
export function search(chatId: string, query: string, opts: { limit?: number } = {}): Passage[] {
  const skipFrom = readJson<IndexState>(stateFile(chatId))?.skipFrom ?? Infinity;
  const all = loadPassages(chatId).filter((p) => p.m < skipFrom);
  const q = [...new Set(tokenize(query))];
  if (!q.length || !all.length) return [];
  const docs = all.map((p) => tokenize(p.text));
  const avg = docs.reduce((n, d) => n + d.length, 0) / docs.length || 1;
  const df = new Map<string, number>();
  for (const d of docs) for (const w of new Set(d)) df.set(w, (df.get(w) ?? 0) + 1);
  const k1 = 1.4, b = 0.75, N = docs.length;
  const scored = docs.map((d, i) => {
    const tf = new Map<string, number>();
    for (const w of d) tf.set(w, (tf.get(w) ?? 0) + 1);
    let s = 0;
    for (const w of q) {
      const f = tf.get(w);
      if (!f) continue;
      const idf = Math.log(1 + (N - (df.get(w) ?? 0) + 0.5) / ((df.get(w) ?? 0) + 0.5));
      s += idf * ((f * (k1 + 1)) / (f + k1 * (1 - b + (b * d.length) / avg)));
    }
    return { p: all[i], s };
  });
  return scored.filter((x) => x.s > 0).sort((a, b2) => b2.s - a.s).slice(0, opts.limit ?? 4).map((x) => x.p);
}

export function formatResults(results: Passage[]): string {
  if (!results.length) return "status: no_results\nnext: Nothing earlier in this chat matches. Try other words, or tell the user you could not find it.";
  const lines = ["status: ok"];
  let used = 0;
  for (const p of results.sort((a, b) => a.m - b.m)) {
    const who = p.role === "user" ? "the user" : p.role === "assistant" ? "the assistant" : `tool ${p.role.slice(5)}`;
    const entry = `[turn ${p.turn}, ${who}, ${p.ts.slice(0, 16).replace("T", " ")} UTC]\n${clip(p.text, Math.min(1600, RESULT_CHARS - used))}`;
    if (used + entry.length > RESULT_CHARS) break;
    lines.push(entry);
    used += entry.length;
  }
  lines.push("next: These are the original words from earlier in this chat. Cite the turn when you rely on them, for example as a finding's source (\"turn 12\").");
  return lines.join("\n\n");
}

// The first message index whose content is still within the verbatim window, counting back
// from the newest message. Passages from there on are in the model's view already.
export function verbatimStart(history: Chat, windowTokens: number): number {
  const msgs = history.getMessagesArray();
  let chars = 0;
  for (let i = msgs.length - 1; i >= 0; i--) {
    chars += msgs[i].getText().length + msgs[i].getToolCallResults().reduce((n, r) => n + r.content.length, 0);
    if (chars / 4 > windowTokens) return i + 1;
  }
  return 0;
}