docx

import { tool, text, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import * as fs from "node:fs/promises";
import * as os from "node:os";
import * as path from "node:path";
import { configSchematics } from "./configSchematics";
import {
  canonicalizeRoots,
  resolveSafe,
  isInsideBackupDir,
  BACKUP_DIR_NAME,
  PathError,
} from "./pathGuard";
import { readDocx } from "./docxRead";
import { writeDocx } from "./docxWrite";
import { replaceInDocx } from "./docxReplace";
import { appendAudit } from "./auditLog";

const LOG_PREFIX = "[docx]";
const DEFAULT_ROOT = path.join(os.homedir(), "Documents");

function rootsHint(roots: string[]): string {
  return roots.length ? roots.join(", ") : "(none configured — enable the default root or set 'Additional allowed root paths' first)";
}

async function backupIfExists(abs: string, root: string): Promise<string | null> {
  let exists = false;
  let isFile = false;
  try {
    const st = await fs.stat(abs);
    exists = true;
    isFile = st.isFile();
  } catch {
    return null;
  }
  if (!exists) return null;
  if (!isFile) {
    throw new Error(`"${abs}" exists but is not a regular file.`);
  }
  const rel = path.relative(root, abs);
  const ts = new Date().toISOString().replace(/[:.]/g, "-");
  const backupPath = path.join(root, BACKUP_DIR_NAME, `${rel}.${ts}.bak`);
  await fs.mkdir(path.dirname(backupPath), { recursive: true });
  await fs.copyFile(abs, backupPath);
  return backupPath;
}

export async function toolsProvider(ctl: ToolsProviderController) {
  const config = ctl.getPluginConfig(configSchematics);
  const useDefaultRoot = config.get("useDefaultRoot");
  const rawRoots = config.get("allowedPaths");
  const combinedRoots = useDefaultRoot ? [DEFAULT_ROOT, ...rawRoots] : rawRoots;
  const roots = await canonicalizeRoots(combinedRoots);
  const fontFamily = config.get("defaultFontFamily");
  const fontSizePt = config.get("defaultFontSizePt");
  const pageMargin = config.get("pageMargin") as "normal" | "narrow" | "wide";
  const preserveListStyles = config.get("preserveListStyles");
  const maxFileSizeMb = config.get("maxFileSizeMb");
  const verboseLogging = config.get("verboseLogging");

  const rootsList = rootsHint(roots);

  const readDocxToMarkdown = tool({
    name: "read_docx_to_markdown",
    description: text`
      Read a Microsoft Word .docx file and return its content as Markdown.

      Use this when the user asks you to:
        • read / open / résumer / fiche / lire / analyser un .docx
        • extract text from a Word document
        • feed a contract, CV, report, or meeting notes into another tool

      How it works (so you set expectations correctly):
        • Purely structural — reads the .docx XML (no vision model, no OCR).
        • Fast: a 30-page document takes well under a second.
        • Headings, bold/italic, lists, blockquotes, links, basic tables
          are preserved. Images are dropped in this iteration.
        • Everything stays local — no network call leaves the machine.

      Allowed root directories: ${rootsList}.
      Files larger than ${maxFileSizeMb} MB are refused.
    `,
    parameters: {
      path: z.string().min(1).describe("Absolute path to the .docx file."),
      preserve_styles: z.boolean().optional().describe(
        "Keep heading levels, bold/italic, blockquote styles. Default true.",
      ),
      include_metadata: z.boolean().optional().describe(
        "Prepend a YAML frontmatter with title/author/modified from the .docx core properties. Default false.",
      ),
    },
    implementation: async (args, ctx) => {
      try {
        const { abs } = await resolveSafe(args.path, roots);
        if (path.extname(abs).toLowerCase() !== ".docx") {
          return { error: `Expected a .docx file, got: ${abs}` };
        }
        const stat = await fs.stat(abs);
        if (!stat.isFile()) return { error: `"${abs}" is not a regular file.` };
        const sizeMb = stat.size / (1024 * 1024);
        if (sizeMb > maxFileSizeMb) {
          return {
            error: `File is ${sizeMb.toFixed(1)} MB, exceeds the ${maxFileSizeMb} MB cap. Raise 'Max file size' in plugin settings.`,
          };
        }

        ctx.status(`Reading ${path.basename(abs)} (${sizeMb.toFixed(1)} MB)…`);
        const t0 = Date.now();
        const result = await readDocx(abs, {
          preserveStyles: args.preserve_styles ?? true,
          includeMetadata: args.include_metadata ?? false,
        });
        const ms = Date.now() - t0;
        ctx.status(`Done — ${result.source_chars} chars in ${ms} ms.`);
        if (verboseLogging) {
          console.log(`${LOG_PREFIX} read ${abs}: ${result.source_chars} chars in ${ms} ms`);
        }
        return result;
      } catch (e) {
        if (e instanceof PathError) return { error: e.message };
        return { error: e instanceof Error ? e.message : String(e) };
      }
    },
  });

  const writeMarkdownToDocx = tool({
    name: "write_markdown_to_docx",
    description: text`
      Convert Markdown to a Microsoft Word .docx file and write it.

      **You DO have filesystem write access through this tool.** Call it once
      with \`path\` and \`markdown\` to produce the .docx. Do NOT refuse or
      claim "I can't write files locally" — you can.

      Existing files are backed up to \`<root>/${BACKUP_DIR_NAME}/\` before
      being overwritten — the previous version is recoverable. Every write is
      logged to \`<root>/${BACKUP_DIR_NAME}/log.md\`.

      Mapping (i1):
        • # ## ### → Heading 1/2/3
        • **bold** *italic* \`code\` → bold / italic / monospace runs
        • - / 1. lists → bulleted / numbered Word lists (sub-levels supported)
        • > quote → indented italic paragraph
        • \`\`\`code\`\`\` → monospace block
        • [link](url) → Word hyperlink
        • Images and complex tables are NOT supported in i1.

      Allowed root directories: ${rootsList}.
    `,
    parameters: {
      path: z.string().min(1).describe("Absolute path of the .docx to create."),
      markdown: z.string().min(1).describe("Markdown content to convert."),
      title: z.string().optional().describe("Document title set in .docx properties."),
    },
    implementation: async (args, ctx) => {
      try {
        const { abs, root } = await resolveSafe(args.path, roots);
        if (path.extname(abs).toLowerCase() !== ".docx") {
          return { error: `Output path must end in .docx, got: ${abs}` };
        }
        if (isInsideBackupDir(abs, root)) {
          return { error: `Refusing to write inside the backup directory (${BACKUP_DIR_NAME}).` };
        }

        ctx.status(`Converting markdown to ${path.basename(abs)}…`);
        const backupPath = await backupIfExists(abs, root);
        await fs.mkdir(path.dirname(abs), { recursive: true });

        const t0 = Date.now();
        const result = await writeDocx(abs, args.markdown, {
          fontFamily,
          fontSizePt,
          pageMargin,
          preserveListStyles,
          title: args.title,
        });
        const ms = Date.now() - t0;
        ctx.status(
          `Wrote ${result.bytes} bytes, ${result.paragraphs} paragraph(s)${backupPath ? " (backup saved)" : ""} in ${ms} ms.`,
        );
        if (verboseLogging) {
          console.log(
            `${LOG_PREFIX} wrote ${abs}: ${result.bytes} bytes, ${result.paragraphs} paragraphs in ${ms} ms${backupPath ? ` (backup: ${backupPath})` : ""}`,
          );
        }
        await appendAudit(
          root,
          backupPath ? "docx-update" : "docx-write",
          path.relative(root, abs),
          `${result.bytes} bytes, ${result.paragraphs} paragraph(s)${backupPath ? " (backup saved)" : ""}`,
        );
        for (const w of result.warnings) ctx.warn(w);
        return {
          written: {
            path: abs,
            bytes: result.bytes,
            paragraphs: result.paragraphs,
          },
          backup_path: backupPath,
          warnings: result.warnings,
        };
      } catch (e) {
        if (e instanceof PathError) return { error: e.message };
        return { error: e instanceof Error ? e.message : String(e) };
      }
    },
  });

  const replaceTextInDocx = tool({
    name: "replace_text_in_docx",
    description: text`
      Replace one or more text strings inside a .docx file, **preserving the
      original Word formatting bit-for-bit** (fonts, colors, layout, headers,
      footers, comments, etc.). Use this when the user wants to edit specific
      text in a Word document without re-rendering it through markdown.

      Typical use cases:
        • Anonymise a contract: replace names, emails, account numbers.
        • Rename a project / company across a long document.
        • Apply a list of corrections without losing the original styling.

      How it differs from \`write_markdown_to_docx\`:
        • \`write_markdown_to_docx\` rebuilds a fresh .docx from markdown —
          structure is preserved but custom Word styling is replaced by this
          plugin's defaults.
        • \`replace_text_in_docx\` opens the original .docx, edits text in
          place inside the XML, and saves a copy. Original styling is fully
          preserved — only the matched substrings change.

      Caveats (read these to set expectations):
        • Replacement is plain string substitution, not regex.
        • Word sometimes splits text across multiple "runs" (because of
          spell-check marks or inline formatting changes). When a search
          string lives across such a split, this tool cannot replace it
          atomically and emits a warning saying so. In that case, fall back
          to the markdown round-trip pipeline (read → edit → write).
        • Replacement scope: \`word/document.xml\`, all \`header*.xml\`,
          \`footer*.xml\`, \`footnotes.xml\`, \`endnotes.xml\`, \`comments.xml\`.
          Custom style names and field codes are not touched.

      Existing files at the output path are backed up to
      \`<root>/${BACKUP_DIR_NAME}/\` before being overwritten. Every write is
      logged to \`<root>/${BACKUP_DIR_NAME}/log.md\`.

      Allowed root directories: ${rootsList}.
    `,
    parameters: {
      path: z.string().min(1).describe("Absolute path to the source .docx."),
      replacements: z
        .array(
          z.object({
            find: z.string().min(1).describe("Exact substring to search for."),
            replace: z.string().describe("Replacement (can be empty to delete)."),
          }),
        )
        .min(1)
        .describe("List of (find, replace) pairs applied in order."),
      output_path: z
        .string()
        .optional()
        .describe(
          "Where to write the modified .docx. Defaults to the input path (in-place edit, with automatic backup).",
        ),
      case_sensitive: z
        .boolean()
        .optional()
        .describe("Match case-sensitively. Default true."),
    },
    implementation: async (args, ctx) => {
      try {
        const { abs: inAbs } = await resolveSafe(args.path, roots);
        if (path.extname(inAbs).toLowerCase() !== ".docx") {
          return { error: `Expected a .docx file, got: ${inAbs}` };
        }
        const stat = await fs.stat(inAbs);
        if (!stat.isFile()) return { error: `"${inAbs}" is not a regular file.` };
        const sizeMb = stat.size / (1024 * 1024);
        if (sizeMb > maxFileSizeMb) {
          return {
            error: `Source is ${sizeMb.toFixed(1)} MB, exceeds the ${maxFileSizeMb} MB cap.`,
          };
        }

        const outRaw = args.output_path ?? args.path;
        const { abs: outAbs, root: outRoot } = await resolveSafe(outRaw, roots);
        if (path.extname(outAbs).toLowerCase() !== ".docx") {
          return { error: `Output path must end in .docx, got: ${outAbs}` };
        }
        if (isInsideBackupDir(outAbs, outRoot)) {
          return { error: `Refusing to write inside the backup directory (${BACKUP_DIR_NAME}).` };
        }

        ctx.status(
          `Replacing ${args.replacements.length} pattern(s) in ${path.basename(inAbs)}…`,
        );
        const backupPath = await backupIfExists(outAbs, outRoot);
        await fs.mkdir(path.dirname(outAbs), { recursive: true });

        const t0 = Date.now();
        const result = await replaceInDocx(
          inAbs,
          outAbs,
          args.replacements,
          { caseSensitive: args.case_sensitive ?? true },
        );
        const ms = Date.now() - t0;
        const total = result.replacements.reduce((s, r) => s + r.count, 0);
        ctx.status(
          `Done — ${total} substitution(s)${backupPath ? " (backup saved)" : ""} in ${ms} ms.`,
        );
        if (verboseLogging) {
          console.log(
            `${LOG_PREFIX} replace ${inAbs} → ${outAbs}: ${total} substitutions in ${ms} ms${backupPath ? ` (backup: ${backupPath})` : ""}`,
          );
        }
        await appendAudit(
          outRoot,
          "docx-replace",
          path.relative(outRoot, outAbs),
          `${total} substitution(s), ${result.bytes} bytes${backupPath ? " (backup saved)" : ""}`,
        );
        for (const w of result.warnings) ctx.warn(w);
        return {
          written: { path: outAbs, bytes: result.bytes },
          replacements: result.replacements,
          backup_path: backupPath,
          warnings: result.warnings,
        };
      } catch (e) {
        if (e instanceof PathError) return { error: e.message };
        return { error: e instanceof Error ? e.message : String(e) };
      }
    },
  });

  return [readDocxToMarkdown, writeMarkdownToDocx, replaceTextInDocx];
}