// Files the user attaches to a chat. LM Studio keeps attachments in its own folder, which
// simpleFiles never opens, and its document reader hands a PDF to the model as plain text, which
// loses a form's fields. So each attached document is copied into an attachments folder that
// simpleFiles treats as granted, and the model is told where the copy is. A PDF with fillable
// fields is also taken off the message, so the document reader does not paste its text in as
// well; the model reads it with read_form instead. Images stay with the message for the model to see.
//
// LM Studio runs the document reader (lmstudio/rag-v1) before simpleFiles when it comes first in
// the chat's plugin list. It then takes the documents off the message itself, and simpleFiles
// sees only the text it wrote. In that case the files are found in LM Studio's user-files
// folder by the names the reader wrote, and a form's pasted text is taken back out.

import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import type { ChatMessage, FileHandle, LMStudioClient } from "@lmstudio/sdk";
import { sha256 } from "./store";

export const defaultAttachmentsFolder = () => path.join(os.homedir(), "Documents", "simpleFiles attachments");

// The folder from the setting: empty means the default, "off" turns copying off.
export function attachmentsFolder(setting: string | undefined): string | null {
  const v = (setting ?? "").trim().replace(/^["']|["']$/g, "");
  if (/^off$/i.test(v)) return null;
  return v || defaultAttachmentsFolder();
}

// A name in the folder for the file: its own name, or the same file already copied there, or
// the name with (2), (3)... when a different file has it.
function placeFor(folder: string, name: string, bytes: Buffer): { file: string; existed: boolean } {
  const ext = path.extname(name);
  const base = path.basename(name, ext).replace(/[<>:"/\\|?*\u0000-\u001f]/g, "_") || "attachment";
  const hash = sha256(bytes);
  for (let n = 1; n < 1000; n++) {
    const file = path.join(folder, n === 1 ? `${base}${ext}` : `${base} (${n})${ext}`);
    if (!fs.existsSync(file)) return { file, existed: false };
    try {
      if (sha256(fs.readFileSync(file)) === hash) return { file, existed: true };
    } catch {
      // unreadable; try the next name
    }
  }
  return { file: path.join(folder, `${base} (${Date.now()})${ext}`), existed: false };
}

async function formFieldCount(bytes: Buffer): Promise<number> {
  try {
    const { PDFDocument } = await import("pdf-lib");
    const doc = await PDFDocument.load(bytes, { ignoreEncryption: true, updateMetadata: false });
    // Counted the way read_form numbers them: each box of a choice is its own field.
    return doc.getForm().getFields().reduce((n, f) => n + f.acroField.getWidgets().length, 0);
  } catch {
    return 0;
  }
}

// One attached document: from the message itself, or found in LM Studio's user-files folder.
interface Attached {
  identifier: string;
  name: string;
  type: string;
  path: () => Promise<string>;
}

// What lmstudio/rag-v1 writes in place of the user's message.
const FULL_START = "This is a Enriched Context Generation scenario.";
const FULL_QUERY = "Based on the content above, please provide a response to the user query.\n\nUser query: ";
const CITED = /^(The following citations were found in the files provided by the user:|Important: No citations were found in the user files for the user query\.)/;
// How long after a file is attached it may still be the one this message carries.
const RECENT_MS = 60 * 60_000;

// LM Studio's folder of attached files: beside its other folders under .lmstudio, found from
// where this plugin runs, or in the home folder.
export function userFilesFolder(): string | null {
  const tries: string[] = [];
  if (process.env.SIMPLEFILES_USER_FILES) tries.push(process.env.SIMPLEFILES_USER_FILES);
  let dir = __dirname;
  for (let i = 0; i < 8; i++) {
    if (path.basename(dir).toLowerCase() === ".lmstudio") tries.push(path.join(dir, "user-files"));
    const up = path.dirname(dir);
    if (up === dir) break;
    dir = up;
  }
  tries.push(path.join(os.homedir(), ".lmstudio", "user-files"));
  return tries.find((d) => fs.existsSync(d)) ?? null;
}

interface Meta {
  fileIdentifier?: string;
  originalName?: string;
  type?: string;
}

// Files attached in the last hour, newest first. A file's name in the folder starts with the
// time it was attached, in milliseconds.
function recentUserFiles(folder: string): Attached[] {
  let names: string[];
  try {
    names = fs.readdirSync(folder).filter((n) => n.endsWith(".metadata.json"));
  } catch {
    return [];
  }
  const out: { at: number; file: Attached }[] = [];
  for (const n of names) {
    const at = Number(/^(\d{12,})/.exec(n)?.[1]);
    if (!at || Date.now() - at > RECENT_MS || at > Date.now() + 60_000) continue;
    let meta: Meta;
    try {
      meta = JSON.parse(fs.readFileSync(path.join(folder, n), "utf8")) as Meta;
    } catch {
      continue;
    }
    const id = meta.fileIdentifier ?? n.replace(/\.metadata\.json$/, "");
    const file = path.join(folder, id);
    if (!fs.existsSync(file)) continue;
    const type = meta.type ?? "";
    if (type === "image" || type.startsWith("image/")) continue;
    out.push({ at, file: { identifier: id, name: meta.originalName ?? id, type, path: async () => file } });
  }
  return out.sort((a, b) => b.at - a.at).map((o) => o.file);
}

// The names rag-v1 wrote when it pasted documents in whole.
function pastedNames(text: string): string[] {
  const names: string[] = [];
  for (const m of text.matchAll(/\n\*\* (.+?) full content \*\*\n/g)) if (text.includes(`** end of ${m[1]} **`)) names.push(m[1]);
  return names;
}

// When rag-v1 already took the documents, finds them by the names it wrote. When it quoted
// pieces instead, it does not name them, so the recent files the chat has not seen before are taken.
async function filesTheReaderTook(text: string, history: () => Promise<Set<string>>): Promise<Attached[]> {
  const whole = text.startsWith(FULL_START) && text.includes(FULL_QUERY);
  if (!whole && !CITED.test(text)) return [];
  const folder = userFilesFolder();
  if (!folder) return [];
  const recent = recentUserFiles(folder);
  if (whole) {
    const found: Attached[] = [];
    for (const name of pastedNames(text)) {
      const f = recent.find((r) => r.name === name);
      if (f && !found.includes(f)) found.push(f);
    }
    return found;
  }
  let seen: Set<string>;
  try {
    seen = await history();
  } catch {
    return [];
  }
  return recent.filter((r) => !seen.has(r.identifier));
}

// Takes a form's pasted text back out of what rag-v1 wrote; with nothing else pasted, leaves
// just the user's own words.
function unpaste(text: string, names: string[]): string {
  let out = text;
  for (const name of names) {
    const start = out.indexOf(`\n\n** ${name} full content **\n\n`);
    const endMark = `** end of ${name} **\n\n`;
    const end = out.indexOf(endMark, start);
    if (start < 0 || end < 0) continue;
    out = out.slice(0, start) + out.slice(end + endMark.length);
  }
  if (!pastedNames(out).length && out.startsWith(FULL_START)) {
    const at = out.indexOf(FULL_QUERY);
    if (at >= 0) return out.slice(at + FULL_QUERY.length);
  }
  return out;
}

export interface Taken {
  message: ChatMessage;
  note: string;
}

// Copies the message's documents into the folder. Returns the message to pass on (forms taken
// off it) and a note for the model, or null when nothing was attached or copying is off.
// history gives the identifiers of files already attached earlier in the chat.
export async function takeAttachments(
  client: LMStudioClient,
  userMessage: ChatMessage,
  folder: string | null,
  history: () => Promise<Set<string>> = async () => new Set()
): Promise<Taken | null> {
  if (!folder) return null;
  let docs: Attached[] = [];
  let onMessage = false;
  try {
    if (userMessage.hasFiles()) {
      onMessage = true;
      docs = userMessage
        .getFiles(client)
        .filter((f: FileHandle) => f.type !== "image")
        .map((f: FileHandle) => ({ identifier: f.identifier, name: f.name, type: f.type, path: () => f.getFilePath() }));
    } else {
      docs = await filesTheReaderTook(userMessage.getText(), history);
    }
  } catch {
    return null;
  }
  if (!docs.length) return null;
  const lines: string[] = [];
  const forms = new Set<string>();
  const formNames: string[] = [];
  for (const f of docs) {
    let from: string;
    let bytes: Buffer;
    try {
      from = await f.path();
      bytes = fs.readFileSync(from);
    } catch (e) {
      lines.push(`"${f.name}" could not be copied (${(e as Error).message}).`);
      continue;
    }
    let place: { file: string; existed: boolean };
    try {
      fs.mkdirSync(folder, { recursive: true });
      place = placeFor(folder, f.name, bytes);
      if (!place.existed) fs.copyFileSync(from, place.file);
    } catch (e) {
      lines.push(`"${f.name}" could not be copied into ${folder} (${(e as Error).message}).`);
      continue;
    }
    const isPdf = f.type === "application/pdf" || /\.pdf$/i.test(f.name);
    const fields = isPdf ? await formFieldCount(bytes) : 0;
    if (fields > 0) {
      forms.add(f.identifier);
      formNames.push(f.name);
      lines.push(`"${f.name}" is a fillable form with ${fields} fields; a copy is at ${place.file}. Read it with read_form and fill it with fill_form.`);
    } else if (isPdf || f.type === "application/word") {
      lines.push(`"${f.name}": a copy is at ${place.file}. Its text, if any, is in the chat as LM Studio read it.`);
    } else {
      lines.push(`"${f.name}": a copy is at ${place.file}; read it with read_file.`);
    }
  }
  const message = userMessage.asMutableCopy();
  if (forms.size && onMessage) {
    try {
      message.consumeFiles(client, (f) => forms.has(f.identifier));
    } catch {
      // the document reader may paste the form's text too; the copy still works
    }
  } else if (formNames.length) {
    const text = message.getText();
    const plain = unpaste(text, formNames);
    if (plain !== text) message.replaceText(plain);
  }
  return { message, note: `[simpleFiles: the user attached ${docs.length === 1 ? "a file" : `${docs.length} files`}. ${lines.join(" ")}]` };
}

// Files the user attaches to a chat. LM Studio keeps attachments in its own folder, which
// simpleFiles never opens, and its document reader hands a PDF to the model as plain text, which
// loses a form's fields. So each attached document is copied into an attachments folder that
// simpleFiles treats as granted, and the model is told where the copy is. A PDF with fillable
// fields is also taken off the message, so the document reader does not paste its text in as
// well; the model reads it with read_form instead. Images stay with the message for the model to see.
//
// LM Studio runs the document reader (lmstudio/rag-v1) before simpleFiles when it comes first in
// the chat's plugin list. It then takes the documents off the message itself, and simpleFiles
// sees only the text it wrote. In that case the files are found in LM Studio's user-files
// folder by the names the reader wrote, and a form's pasted text is taken back out.

import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import type { ChatMessage, FileHandle, LMStudioClient } from "@lmstudio/sdk";
import { sha256 } from "./store";

export const defaultAttachmentsFolder = () => path.join(os.homedir(), "Documents", "simpleFiles attachments");

// The folder from the setting: empty means the default, "off" turns copying off.
export function attachmentsFolder(setting: string | undefined): string | null {
  const v = (setting ?? "").trim().replace(/^["']|["']$/g, "");
  if (/^off$/i.test(v)) return null;
  return v || defaultAttachmentsFolder();
}

// A name in the folder for the file: its own name, or the same file already copied there, or
// the name with (2), (3)... when a different file has it.
function placeFor(folder: string, name: string, bytes: Buffer): { file: string; existed: boolean } {
  const ext = path.extname(name);
  const base = path.basename(name, ext).replace(/[<>:"/\\|?*\u0000-\u001f]/g, "_") || "attachment";
  const hash = sha256(bytes);
  for (let n = 1; n < 1000; n++) {
    const file = path.join(folder, n === 1 ? `${base}${ext}` : `${base} (${n})${ext}`);
    if (!fs.existsSync(file)) return { file, existed: false };
    try {
      if (sha256(fs.readFileSync(file)) === hash) return { file, existed: true };
    } catch {
      // unreadable; try the next name
    }
  }
  return { file: path.join(folder, `${base} (${Date.now()})${ext}`), existed: false };
}

async function formFieldCount(bytes: Buffer): Promise<number> {
  try {
    const { PDFDocument } = await import("pdf-lib");
    const doc = await PDFDocument.load(bytes, { ignoreEncryption: true, updateMetadata: false });
    // Counted the way read_form numbers them: each box of a choice is its own field.
    return doc.getForm().getFields().reduce((n, f) => n + f.acroField.getWidgets().length, 0);
  } catch {
    return 0;
  }
}

// One attached document: from the message itself, or found in LM Studio's user-files folder.
interface Attached {
  identifier: string;
  name: string;
  type: string;
  path: () => Promise<string>;
}

// What lmstudio/rag-v1 writes in place of the user's message.
const FULL_START = "This is a Enriched Context Generation scenario.";
const FULL_QUERY = "Based on the content above, please provide a response to the user query.\n\nUser query: ";
const CITED = /^(The following citations were found in the files provided by the user:|Important: No citations were found in the user files for the user query\.)/;
// How long after a file is attached it may still be the one this message carries.
const RECENT_MS = 60 * 60_000;

// LM Studio's folder of attached files: beside its other folders under .lmstudio, found from
// where this plugin runs, or in the home folder.
export function userFilesFolder(): string | null {
  const tries: string[] = [];
  if (process.env.SIMPLEFILES_USER_FILES) tries.push(process.env.SIMPLEFILES_USER_FILES);
  let dir = __dirname;
  for (let i = 0; i < 8; i++) {
    if (path.basename(dir).toLowerCase() === ".lmstudio") tries.push(path.join(dir, "user-files"));
    const up = path.dirname(dir);
    if (up === dir) break;
    dir = up;
  }
  tries.push(path.join(os.homedir(), ".lmstudio", "user-files"));
  return tries.find((d) => fs.existsSync(d)) ?? null;
}

interface Meta {
  fileIdentifier?: string;
  originalName?: string;
  type?: string;
}

// Files attached in the last hour, newest first. A file's name in the folder starts with the
// time it was attached, in milliseconds.
function recentUserFiles(folder: string): Attached[] {
  let names: string[];
  try {
    names = fs.readdirSync(folder).filter((n) => n.endsWith(".metadata.json"));
  } catch {
    return [];
  }
  const out: { at: number; file: Attached }[] = [];
  for (const n of names) {
    const at = Number(/^(\d{12,})/.exec(n)?.[1]);
    if (!at || Date.now() - at > RECENT_MS || at > Date.now() + 60_000) continue;
    let meta: Meta;
    try {
      meta = JSON.parse(fs.readFileSync(path.join(folder, n), "utf8")) as Meta;
    } catch {
      continue;
    }
    const id = meta.fileIdentifier ?? n.replace(/\.metadata\.json$/, "");
    const file = path.join(folder, id);
    if (!fs.existsSync(file)) continue;
    const type = meta.type ?? "";
    if (type === "image" || type.startsWith("image/")) continue;
    out.push({ at, file: { identifier: id, name: meta.originalName ?? id, type, path: async () => file } });
  }
  return out.sort((a, b) => b.at - a.at).map((o) => o.file);
}

// The names rag-v1 wrote when it pasted documents in whole.
function pastedNames(text: string): string[] {
  const names: string[] = [];
  for (const m of text.matchAll(/\n\*\* (.+?) full content \*\*\n/g)) if (text.includes(`** end of ${m[1]} **`)) names.push(m[1]);
  return names;
}

// When rag-v1 already took the documents, finds them by the names it wrote. When it quoted
// pieces instead, it does not name them, so the recent files the chat has not seen before are taken.
async function filesTheReaderTook(text: string, history: () => Promise<Set<string>>): Promise<Attached[]> {
  const whole = text.startsWith(FULL_START) && text.includes(FULL_QUERY);
  if (!whole && !CITED.test(text)) return [];
  const folder = userFilesFolder();
  if (!folder) return [];
  const recent = recentUserFiles(folder);
  if (whole) {
    const found: Attached[] = [];
    for (const name of pastedNames(text)) {
      const f = recent.find((r) => r.name === name);
      if (f && !found.includes(f)) found.push(f);
    }
    return found;
  }
  let seen: Set<string>;
  try {
    seen = await history();
  } catch {
    return [];
  }
  return recent.filter((r) => !seen.has(r.identifier));
}

// Takes a form's pasted text back out of what rag-v1 wrote; with nothing else pasted, leaves
// just the user's own words.
function unpaste(text: string, names: string[]): string {
  let out = text;
  for (const name of names) {
    const start = out.indexOf(`\n\n** ${name} full content **\n\n`);
    const endMark = `** end of ${name} **\n\n`;
    const end = out.indexOf(endMark, start);
    if (start < 0 || end < 0) continue;
    out = out.slice(0, start) + out.slice(end + endMark.length);
  }
  if (!pastedNames(out).length && out.startsWith(FULL_START)) {
    const at = out.indexOf(FULL_QUERY);
    if (at >= 0) return out.slice(at + FULL_QUERY.length);
  }
  return out;
}

export interface Taken {
  message: ChatMessage;
  note: string;
}

// Copies the message's documents into the folder. Returns the message to pass on (forms taken
// off it) and a note for the model, or null when nothing was attached or copying is off.
// history gives the identifiers of files already attached earlier in the chat.
export async function takeAttachments(
  client: LMStudioClient,
  userMessage: ChatMessage,
  folder: string | null,
  history: () => Promise<Set<string>> = async () => new Set()
): Promise<Taken | null> {
  if (!folder) return null;
  let docs: Attached[] = [];
  let onMessage = false;
  try {
    if (userMessage.hasFiles()) {
      onMessage = true;
      docs = userMessage
        .getFiles(client)
        .filter((f: FileHandle) => f.type !== "image")
        .map((f: FileHandle) => ({ identifier: f.identifier, name: f.name, type: f.type, path: () => f.getFilePath() }));
    } else {
      docs = await filesTheReaderTook(userMessage.getText(), history);
    }
  } catch {
    return null;
  }
  if (!docs.length) return null;
  const lines: string[] = [];
  const forms = new Set<string>();
  const formNames: string[] = [];
  for (const f of docs) {
    let from: string;
    let bytes: Buffer;
    try {
      from = await f.path();
      bytes = fs.readFileSync(from);
    } catch (e) {
      lines.push(`"${f.name}" could not be copied (${(e as Error).message}).`);
      continue;
    }
    let place: { file: string; existed: boolean };
    try {
      fs.mkdirSync(folder, { recursive: true });
      place = placeFor(folder, f.name, bytes);
      if (!place.existed) fs.copyFileSync(from, place.file);
    } catch (e) {
      lines.push(`"${f.name}" could not be copied into ${folder} (${(e as Error).message}).`);
      continue;
    }
    const isPdf = f.type === "application/pdf" || /\.pdf$/i.test(f.name);
    const fields = isPdf ? await formFieldCount(bytes) : 0;
    if (fields > 0) {
      forms.add(f.identifier);
      formNames.push(f.name);
      lines.push(`"${f.name}" is a fillable form with ${fields} fields; a copy is at ${place.file}. Read it with read_form and fill it with fill_form.`);
    } else if (isPdf || f.type === "application/word") {
      lines.push(`"${f.name}": a copy is at ${place.file}. Its text, if any, is in the chat as LM Studio read it.`);
    } else {
      lines.push(`"${f.name}": a copy is at ${place.file}; read it with read_file.`);
    }
  }
  const message = userMessage.asMutableCopy();
  if (forms.size && onMessage) {
    try {
      message.consumeFiles(client, (f) => forms.has(f.identifier));
    } catch {
      // the document reader may paste the form's text too; the copy still works
    }
  } else if (formNames.length) {
    const text = message.getText();
    const plain = unpaste(text, formNames);
    if (plain !== text) message.replaceText(plain);
  }
  return { message, note: `[simpleFiles: the user attached ${docs.length === 1 ? "a file" : `${docs.length} files`}. ${lines.join(" ")}]` };
}