// What the model has been shown from this computer, shared between plugins so that one which
// sends things out (simpleSearch) can keep personal details from leaving.
//
// The same file is in simpleSearch, simpleFiles and simpleMemory; keep the copies identical,
// since every side must turn text into the same fingerprints.
//
// What is recorded depends on where the text came from:
// - Any local text (a file, the memory file): its identifiers. These are strings that point at
//   one person or thing and are almost never typed by chance: email addresses, phone, account
//   and order numbers (and any seven digits in a row of them), and keys, serials and IDs.
// - The memory file only: also its phrases, runs of six uncommon words. Memory is written about
//   the user, so a long run copied from it is personal. Files are mostly code, notes and
//   documents whose ordinary wording turns up everywhere, so their phrases are not recorded.
// Names and places cannot be told apart from ordinary words by their shape; the user lists the
// ones that matter on simpleSearch's Never send list.
//
// Fingerprints are short hashes, so no second copy of the text is kept.
//
// A plugin that shows the user's files also puts a line starting with "privacy: local" in its
// output and records the time, for simpleSearch's optional ask-first mode.

import * as crypto from "node:crypto";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";

export const MARK_PATTERN = /^privacy: local\b(?:, from ([\w-]+))?/gm;
export const markLine = (source: string) => `privacy: local, from ${source}`;

const PHRASE = 6;
const DIGIT_WINDOW = 7;
const MAX_FILE_BYTES = 2 * 1024 * 1024;
// Fingerprints from earlier versions that are no longer made; ignored when read.
const RETIRED = /^r/;

export function sharedDir(): string {
  return process.env.LOCAL_ONLY_DIR ? path.resolve(process.env.LOCAL_ONLY_DIR) : path.join(os.homedir(), ".lmstudio", "local-only");
}
const shownFile = () => path.join(sharedDir(), "shown.json");
const printsFile = () => path.join(sharedDir(), "fingerprints.txt");

// Words that say nothing about whose text it is; phrases skip them, so adding or dropping one
// does not change a phrase.
const COMMON = new Set(
  (
    "a an and are as at be been but by can could do does did for from had has have he her his how i if in into is it its " +
    "me my no not of on or our she so than that the their them then there these they this those to too us was we were what " +
    "when where which who why will with would you your about after all also any because before between both each few more " +
    "most other over same should some such through under until very while just only own here out up down off again further once"
  ).split(" ")
);

export function words(text: string): string[] {
  return text.normalize("NFKC").toLowerCase().match(/[\p{L}\p{N}]+/gu) ?? [];
}

const hash = (s: string) => crypto.createHash("sha256").update(s).digest("hex").slice(0, 16);

// Email addresses, numbers of six digits or more, and unbroken codes of ten characters or more
// that mix letters and digits. Dates are skipped, and so are codes with a run of five letters
// that reads like a word (hex digits aside), so a model name such as "qwen3-35b-uncensored" is
// left alone.
export function distinctive(text: string): string[] {
  const t = text.normalize("NFKC").toLowerCase();
  const found = new Set<string>();
  for (const m of t.matchAll(/[\w.+-]+@[\w-]+(?:\.[\w-]+)+/g)) found.add(m[0]);
  for (const m of t.matchAll(/\d[\d ().-]{4,}\d/g)) {
    if (/^\d{4}[-/.]\d{1,2}[-/.]\d{1,2}\b/.test(m[0]) || /^\d{1,2}[-/.]\d{1,2}[-/.]\d{2,4}\b/.test(m[0])) continue;
    const digits = m[0].replace(/\D/g, "");
    if (digits.length >= 6) found.add(digits);
    for (let i = 0; i + DIGIT_WINDOW <= digits.length; i++) found.add("#" + digits.slice(i, i + DIGIT_WINDOW));
  }
  for (const m of t.matchAll(/[a-z0-9]{10,}/g)) {
    const c = m[0];
    if ((c.match(/\d/g) ?? []).length >= 2 && (c.match(/[a-z]/g) ?? []).length >= 2 && !(c.match(/[a-z]{5,}/g) ?? []).some((run) => /[g-z]/.test(run))) found.add(c);
  }
  return [...found];
}

export function identifierPrints(text: string): string[] {
  return distinctive(text).map((d) => "t" + hash(d));
}

// Runs of six uncommon words, at least three of them four letters or longer.
export function phrasePrints(text: string): string[] {
  const rare = words(text).filter((w) => !COMMON.has(w));
  const out: string[] = [];
  for (let i = 0; i + PHRASE <= rare.length; i++) {
    const run = rare.slice(i, i + PHRASE);
    if (run.filter((w) => w.length >= 4).length < 3) continue;
    out.push("p" + hash(run.join(" ")));
  }
  return out;
}

// Everything a piece of outgoing text could match.
export function fingerprints(text: string): string[] {
  return [...new Set([...identifierPrints(text), ...phrasePrints(text)])];
}

function record(prints: string[]): void {
  try {
    const known = knownPrints();
    const fresh = [...new Set(prints)].filter((p) => !known.has(p));
    if (!fresh.length) return;
    fs.mkdirSync(sharedDir(), { recursive: true });
    fs.appendFileSync(printsFile(), fresh.join("\n") + "\n");
    trim();
  } catch {
    // Best effort.
  }
}

// Called by a plugin that has just shown the model the user's files: records the time and the
// identifiers in what it showed.
export function recordShown(source: string, text: string): void {
  try {
    fs.mkdirSync(sharedDir(), { recursive: true });
    const shown = readShown();
    shown[source] = Date.now();
    const tmp = `${shownFile()}.${process.pid}.tmp`;
    fs.writeFileSync(tmp, JSON.stringify(shown));
    fs.renameSync(tmp, shownFile());
  } catch {
    // Best effort: the history mark still tells the searching plugin.
  }
  record(identifierPrints(text));
}

// Called with text written about the user, such as the memory file: records its identifiers
// and its phrases.
export function recordFingerprints(text: string): void {
  record([...identifierPrints(text), ...phrasePrints(text)]);
}

// Keeps the fingerprint file from growing without end by dropping its older half.
function trim(): void {
  try {
    if (fs.statSync(printsFile()).size <= MAX_FILE_BYTES) return;
    const lines = fs.readFileSync(printsFile(), "utf8").split("\n").filter((l) => l && !RETIRED.test(l));
    const tmp = `${printsFile()}.${process.pid}.tmp`;
    fs.writeFileSync(tmp, lines.slice(Math.floor(lines.length / 2)).join("\n") + "\n");
    fs.renameSync(tmp, printsFile());
  } catch {
    // leave it
  }
}

// When each source last showed something: source name to milliseconds since 1970.
export function readShown(): Record<string, number> {
  try {
    const v = JSON.parse(fs.readFileSync(shownFile(), "utf8"));
    return v && typeof v === "object" ? v : {};
  } catch {
    return {};
  }
}

let cache: { mtime: number; size: number; set: Set<string> } | null = null;

export function knownPrints(): Set<string> {
  let st: fs.Stats;
  try {
    st = fs.statSync(printsFile());
  } catch {
    return new Set();
  }
  if (!cache || cache.mtime !== st.mtimeMs || cache.size !== st.size) {
    const set = new Set(fs.readFileSync(printsFile(), "utf8").split("\n").filter((l) => l && !RETIRED.test(l)));
    cache = { mtime: st.mtimeMs, size: st.size, set };
  }
  return cache.set;
}

// True when the text repeats an identifier from local text or a phrase from the memory file.
// Fingerprints in except (for example, of what the user typed in the chat) do not count.
export function hasCopiedText(text: string, except: Set<string> = new Set()): boolean {
  const known = knownPrints();
  if (!known.size) return false;
  return fingerprints(text).some((p) => known.has(p) && !except.has(p));
}

// The sources named in "privacy: local" lines in a piece of text.
export function marksIn(text: string): string[] {
  const out: string[] = [];
  for (const m of text.matchAll(MARK_PATTERN)) out.push(m[1] ?? "unknown");
  return out;
}

// What the model has been shown from this computer, shared between plugins so that one which
// sends things out (simpleSearch) can keep personal details from leaving.
//
// The same file is in simpleSearch, simpleFiles and simpleMemory; keep the copies identical,
// since every side must turn text into the same fingerprints.
//
// What is recorded depends on where the text came from:
// - Any local text (a file, the memory file): its identifiers. These are strings that point at
//   one person or thing and are almost never typed by chance: email addresses, phone, account
//   and order numbers (and any seven digits in a row of them), and keys, serials and IDs.
// - The memory file only: also its phrases, runs of six uncommon words. Memory is written about
//   the user, so a long run copied from it is personal. Files are mostly code, notes and
//   documents whose ordinary wording turns up everywhere, so their phrases are not recorded.
// Names and places cannot be told apart from ordinary words by their shape; the user lists the
// ones that matter on simpleSearch's Never send list.
//
// Fingerprints are short hashes, so no second copy of the text is kept.
//
// A plugin that shows the user's files also puts a line starting with "privacy: local" in its
// output and records the time, for simpleSearch's optional ask-first mode.

import * as crypto from "node:crypto";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";

export const MARK_PATTERN = /^privacy: local\b(?:, from ([\w-]+))?/gm;
export const markLine = (source: string) => `privacy: local, from ${source}`;

const PHRASE = 6;
const DIGIT_WINDOW = 7;
const MAX_FILE_BYTES = 2 * 1024 * 1024;
// Fingerprints from earlier versions that are no longer made; ignored when read.
const RETIRED = /^r/;

export function sharedDir(): string {
  return process.env.LOCAL_ONLY_DIR ? path.resolve(process.env.LOCAL_ONLY_DIR) : path.join(os.homedir(), ".lmstudio", "local-only");
}
const shownFile = () => path.join(sharedDir(), "shown.json");
const printsFile = () => path.join(sharedDir(), "fingerprints.txt");

// Words that say nothing about whose text it is; phrases skip them, so adding or dropping one
// does not change a phrase.
const COMMON = new Set(
  (
    "a an and are as at be been but by can could do does did for from had has have he her his how i if in into is it its " +
    "me my no not of on or our she so than that the their them then there these they this those to too us was we were what " +
    "when where which who why will with would you your about after all also any because before between both each few more " +
    "most other over same should some such through under until very while just only own here out up down off again further once"
  ).split(" ")
);

export function words(text: string): string[] {
  return text.normalize("NFKC").toLowerCase().match(/[\p{L}\p{N}]+/gu) ?? [];
}

const hash = (s: string) => crypto.createHash("sha256").update(s).digest("hex").slice(0, 16);

// Email addresses, numbers of six digits or more, and unbroken codes of ten characters or more
// that mix letters and digits. Dates are skipped, and so are codes with a run of five letters
// that reads like a word (hex digits aside), so a model name such as "qwen3-35b-uncensored" is
// left alone.
export function distinctive(text: string): string[] {
  const t = text.normalize("NFKC").toLowerCase();
  const found = new Set<string>();
  for (const m of t.matchAll(/[\w.+-]+@[\w-]+(?:\.[\w-]+)+/g)) found.add(m[0]);
  for (const m of t.matchAll(/\d[\d ().-]{4,}\d/g)) {
    if (/^\d{4}[-/.]\d{1,2}[-/.]\d{1,2}\b/.test(m[0]) || /^\d{1,2}[-/.]\d{1,2}[-/.]\d{2,4}\b/.test(m[0])) continue;
    const digits = m[0].replace(/\D/g, "");
    if (digits.length >= 6) found.add(digits);
    for (let i = 0; i + DIGIT_WINDOW <= digits.length; i++) found.add("#" + digits.slice(i, i + DIGIT_WINDOW));
  }
  for (const m of t.matchAll(/[a-z0-9]{10,}/g)) {
    const c = m[0];
    if ((c.match(/\d/g) ?? []).length >= 2 && (c.match(/[a-z]/g) ?? []).length >= 2 && !(c.match(/[a-z]{5,}/g) ?? []).some((run) => /[g-z]/.test(run))) found.add(c);
  }
  return [...found];
}

export function identifierPrints(text: string): string[] {
  return distinctive(text).map((d) => "t" + hash(d));
}

// Runs of six uncommon words, at least three of them four letters or longer.
export function phrasePrints(text: string): string[] {
  const rare = words(text).filter((w) => !COMMON.has(w));
  const out: string[] = [];
  for (let i = 0; i + PHRASE <= rare.length; i++) {
    const run = rare.slice(i, i + PHRASE);
    if (run.filter((w) => w.length >= 4).length < 3) continue;
    out.push("p" + hash(run.join(" ")));
  }
  return out;
}

// Everything a piece of outgoing text could match.
export function fingerprints(text: string): string[] {
  return [...new Set([...identifierPrints(text), ...phrasePrints(text)])];
}

function record(prints: string[]): void {
  try {
    const known = knownPrints();
    const fresh = [...new Set(prints)].filter((p) => !known.has(p));
    if (!fresh.length) return;
    fs.mkdirSync(sharedDir(), { recursive: true });
    fs.appendFileSync(printsFile(), fresh.join("\n") + "\n");
    trim();
  } catch {
    // Best effort.
  }
}

// Called by a plugin that has just shown the model the user's files: records the time and the
// identifiers in what it showed.
export function recordShown(source: string, text: string): void {
  try {
    fs.mkdirSync(sharedDir(), { recursive: true });
    const shown = readShown();
    shown[source] = Date.now();
    const tmp = `${shownFile()}.${process.pid}.tmp`;
    fs.writeFileSync(tmp, JSON.stringify(shown));
    fs.renameSync(tmp, shownFile());
  } catch {
    // Best effort: the history mark still tells the searching plugin.
  }
  record(identifierPrints(text));
}

// Called with text written about the user, such as the memory file: records its identifiers
// and its phrases.
export function recordFingerprints(text: string): void {
  record([...identifierPrints(text), ...phrasePrints(text)]);
}

// Keeps the fingerprint file from growing without end by dropping its older half.
function trim(): void {
  try {
    if (fs.statSync(printsFile()).size <= MAX_FILE_BYTES) return;
    const lines = fs.readFileSync(printsFile(), "utf8").split("\n").filter((l) => l && !RETIRED.test(l));
    const tmp = `${printsFile()}.${process.pid}.tmp`;
    fs.writeFileSync(tmp, lines.slice(Math.floor(lines.length / 2)).join("\n") + "\n");
    fs.renameSync(tmp, printsFile());
  } catch {
    // leave it
  }
}

// When each source last showed something: source name to milliseconds since 1970.
export function readShown(): Record<string, number> {
  try {
    const v = JSON.parse(fs.readFileSync(shownFile(), "utf8"));
    return v && typeof v === "object" ? v : {};
  } catch {
    return {};
  }
}

let cache: { mtime: number; size: number; set: Set<string> } | null = null;

export function knownPrints(): Set<string> {
  let st: fs.Stats;
  try {
    st = fs.statSync(printsFile());
  } catch {
    return new Set();
  }
  if (!cache || cache.mtime !== st.mtimeMs || cache.size !== st.size) {
    const set = new Set(fs.readFileSync(printsFile(), "utf8").split("\n").filter((l) => l && !RETIRED.test(l)));
    cache = { mtime: st.mtimeMs, size: st.size, set };
  }
  return cache.set;
}

// True when the text repeats an identifier from local text or a phrase from the memory file.
// Fingerprints in except (for example, of what the user typed in the chat) do not count.
export function hasCopiedText(text: string, except: Set<string> = new Set()): boolean {
  const known = knownPrints();
  if (!known.size) return false;
  return fingerprints(text).some((p) => known.has(p) && !except.has(p));
}

// The sources named in "privacy: local" lines in a piece of text.
export function marksIn(text: string): string[] {
  const out: string[] = [];
  for (const m of text.matchAll(MARK_PATTERN)) out.push(m[1] ?? "unknown");
  return out;
}