import { tool, text, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./configSchematics";
import { anonymize } from "./anonymizer";
import { detectWithModel, type PiiModelOptions } from "./piiModel";

export async function toolsProvider(ctl: ToolsProviderController) {
  const config = ctl.getPluginConfig(configSchematics);
  const configCustomTerms = config.get("alwaysRedact");
  const international = config.get("detectInternationalPhones");
  // Select field with values "off" / "on" (was a boolean in <1.1, see
  // configSchematics for migration notes).
  const useModel = config.get("detectPiiWithModel") === "on";
  const piiServiceNodePath = config.get("piiServiceNodePath") ?? "";
  const modelOptions: PiiModelOptions = {
    // Defensive `?? <default>`: LM Studio sometimes omits unmodified config
    // keys from the per-chat snapshot, causing config.get() to return
    // undefined. Falling back to the schema default keeps the ML pipeline
    // behaving as the UI shows.
    detectNames: config.get("modelDetectNames") ?? true,
    detectAddresses: config.get("modelDetectAddresses") ?? false,
    detectDates: config.get("modelDetectDates") ?? false,
    detectIdDocs: config.get("modelDetectIdDocs") ?? false,
    nodeBinaryPath: piiServiceNodePath,
  };

  const anonymizeText = tool({
    name: "anonymize_text",
    description: text`
      Redact personally identifiable information (PII) from a block of text.

      The tool DETECTS AUTOMATICALLY (with proper validation, not just regex):
        • Credit / bank card numbers (Luhn-validated)
        • French NIR / numéro de sécurité sociale (mod-97 checksum)
        • IBAN (mod-97 checksum)
        • French phone numbers (0X XX XX XX XX, +33 X XX XX XX XX, 0033…)
        • Email addresses
        • Passport numbers when explicitly labelled (e.g. "Passeport n° …")
        • French addresses when explicitly labelled (e.g. "Adresse : …")

      When the host operator enables it, an ML model additionally detects
      names, addresses, dates, and ID-document numbers (configurable per
      category). Regex+checksum results stay authoritative on overlaps.

      The tool DOES NOT detect names without the ML option. You (the model)
      should identify any personal names in the text and pass them in
      \`names\`. Pass full names like "Jean Dupont", not just first names, to
      avoid false matches against common words. You can also pass arbitrary
      strings to redact via \`custom_terms\` (e.g. company names, addresses,
      project codenames).

      Replacement uses STABLE TYPED PSEUDONYMS:
        \`[NOM_1]\`, \`[CB_1]\`, \`[NIR_1]\`, \`[IBAN_1]\`, \`[TEL_1]\`,
        \`[EMAIL_1]\`, \`[CUSTOM_1]\`, \`[ADDRESS_1]\`, \`[DATE_1]\`,
        \`[IDDOC_1]\`…
      Same value → same pseudonym throughout the text, so coreference is
      preserved (e.g. "Jean … Jean" both become "[NOM_1]"). The response
      includes a \`mapping\` from pseudonym to original value.

      The response also includes \`ml_status\`:
        • \`"disabled"\` — operator has the ML layer off; regex-only result
        • \`"ok"\` — ML ran successfully (names/dates/addresses included)
        • \`"error"\` — ML was enabled but unreachable; see \`ml_error\` for
          the explicit reason (typically: no system \`node\` binary found).
          In that case the regex result is still returned, but you should
          tell the user that ML detection didn't run and how to fix it
          (install Node.js or set \`piiServiceNodePath\` in plugin settings).

      Typical workflow: read a file with the filesystem plugin, call
      anonymize_text on its content, then write the redacted version back.
    `,
    parameters: {
      text: z.string().min(1).describe("The text to anonymize."),
      names: z.array(z.string().min(1)).optional().describe(
        "Personal names to redact. Pass full names ('Jean Dupont'), not just first names.",
      ),
      custom_terms: z.array(z.string().min(1)).optional().describe(
        "Other strings to redact: company names, addresses, project codenames, etc.",
      ),
      include_mapping: z.boolean().optional().describe(
        "Include the pseudonym → original value mapping in the response. Default true.",
      ),
    },
    implementation: async (args, ctx) => {
      const t0 = Date.now();
      ctx.status(`Anonymizing ${args.text.length} chars…`);

      let modelSpans;
      let mlError: string | undefined;
      if (useModel) {
        try {
          modelSpans = await detectWithModel(args.text, {
            ...modelOptions,
            onProgress: (msg) => ctx.status(msg),
          });
        } catch (err) {
          // Don't fail the redaction if the ML service is down — return the
          // regex-only result with an explicit ml_error in the response.
          mlError = err instanceof Error ? err.message : String(err);
          ctx.status(`ML unavailable: ${mlError.slice(0, 100)}`);
        }
      }

      const result = anonymize(args.text, {
        names: args.names,
        customTerms: args.custom_terms,
        configCustomTerms,
        international,
        modelSpans,
      });
      const totalRedactions = Object.values(result.counts).reduce((s, n) => s + n, 0);
      ctx.status(
        `Done — ${totalRedactions} item(s) redacted in ${Date.now() - t0} ms.`,
      );
      // ml_status tells the LLM what really happened with the ML layer.
      //   - "disabled":       user has ML off in config; regex-only result
      //   - "ok":             ML ran, names/dates/addresses included
      //   - "error":          ML enabled but unreachable — see ml_error
      const ml_status: "disabled" | "ok" | "error" =
        !useModel ? "disabled" : (modelSpans ? "ok" : "error");

      return {
        anonymized: result.anonymized,
        counts: result.counts,
        ...(args.include_mapping === false ? {} : { mapping: result.mapping }),
        ml_status,
        ...(mlError ? { ml_error: mlError } : {}),
      };
    },
  });

  return [anonymizeText];
}

import { tool, text, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./configSchematics";
import { anonymize } from "./anonymizer";
import { detectWithModel, type PiiModelOptions } from "./piiModel";

export async function toolsProvider(ctl: ToolsProviderController) {
  const config = ctl.getPluginConfig(configSchematics);
  const configCustomTerms = config.get("alwaysRedact");
  const international = config.get("detectInternationalPhones");
  // Select field with values "off" / "on" (was a boolean in <1.1, see
  // configSchematics for migration notes).
  const useModel = config.get("detectPiiWithModel") === "on";
  const piiServiceNodePath = config.get("piiServiceNodePath") ?? "";
  const modelOptions: PiiModelOptions = {
    // Defensive `?? <default>`: LM Studio sometimes omits unmodified config
    // keys from the per-chat snapshot, causing config.get() to return
    // undefined. Falling back to the schema default keeps the ML pipeline
    // behaving as the UI shows.
    detectNames: config.get("modelDetectNames") ?? true,
    detectAddresses: config.get("modelDetectAddresses") ?? false,
    detectDates: config.get("modelDetectDates") ?? false,
    detectIdDocs: config.get("modelDetectIdDocs") ?? false,
    nodeBinaryPath: piiServiceNodePath,
  };

  const anonymizeText = tool({
    name: "anonymize_text",
    description: text`
      Redact personally identifiable information (PII) from a block of text.

      The tool DETECTS AUTOMATICALLY (with proper validation, not just regex):
        • Credit / bank card numbers (Luhn-validated)
        • French NIR / numéro de sécurité sociale (mod-97 checksum)
        • IBAN (mod-97 checksum)
        • French phone numbers (0X XX XX XX XX, +33 X XX XX XX XX, 0033…)
        • Email addresses
        • Passport numbers when explicitly labelled (e.g. "Passeport n° …")
        • French addresses when explicitly labelled (e.g. "Adresse : …")

      When the host operator enables it, an ML model additionally detects
      names, addresses, dates, and ID-document numbers (configurable per
      category). Regex+checksum results stay authoritative on overlaps.

      The tool DOES NOT detect names without the ML option. You (the model)
      should identify any personal names in the text and pass them in
      \`names\`. Pass full names like "Jean Dupont", not just first names, to
      avoid false matches against common words. You can also pass arbitrary
      strings to redact via \`custom_terms\` (e.g. company names, addresses,
      project codenames).

      Replacement uses STABLE TYPED PSEUDONYMS:
        \`[NOM_1]\`, \`[CB_1]\`, \`[NIR_1]\`, \`[IBAN_1]\`, \`[TEL_1]\`,
        \`[EMAIL_1]\`, \`[CUSTOM_1]\`, \`[ADDRESS_1]\`, \`[DATE_1]\`,
        \`[IDDOC_1]\`…
      Same value → same pseudonym throughout the text, so coreference is
      preserved (e.g. "Jean … Jean" both become "[NOM_1]"). The response
      includes a \`mapping\` from pseudonym to original value.

      The response also includes \`ml_status\`:
        • \`"disabled"\` — operator has the ML layer off; regex-only result
        • \`"ok"\` — ML ran successfully (names/dates/addresses included)
        • \`"error"\` — ML was enabled but unreachable; see \`ml_error\` for
          the explicit reason (typically: no system \`node\` binary found).
          In that case the regex result is still returned, but you should
          tell the user that ML detection didn't run and how to fix it
          (install Node.js or set \`piiServiceNodePath\` in plugin settings).

      Typical workflow: read a file with the filesystem plugin, call
      anonymize_text on its content, then write the redacted version back.
    `,
    parameters: {
      text: z.string().min(1).describe("The text to anonymize."),
      names: z.array(z.string().min(1)).optional().describe(
        "Personal names to redact. Pass full names ('Jean Dupont'), not just first names.",
      ),
      custom_terms: z.array(z.string().min(1)).optional().describe(
        "Other strings to redact: company names, addresses, project codenames, etc.",
      ),
      include_mapping: z.boolean().optional().describe(
        "Include the pseudonym → original value mapping in the response. Default true.",
      ),
    },
    implementation: async (args, ctx) => {
      const t0 = Date.now();
      ctx.status(`Anonymizing ${args.text.length} chars…`);

      let modelSpans;
      let mlError: string | undefined;
      if (useModel) {
        try {
          modelSpans = await detectWithModel(args.text, {
            ...modelOptions,
            onProgress: (msg) => ctx.status(msg),
          });
        } catch (err) {
          // Don't fail the redaction if the ML service is down — return the
          // regex-only result with an explicit ml_error in the response.
          mlError = err instanceof Error ? err.message : String(err);
          ctx.status(`ML unavailable: ${mlError.slice(0, 100)}`);
        }
      }

      const result = anonymize(args.text, {
        names: args.names,
        customTerms: args.custom_terms,
        configCustomTerms,
        international,
        modelSpans,
      });
      const totalRedactions = Object.values(result.counts).reduce((s, n) => s + n, 0);
      ctx.status(
        `Done — ${totalRedactions} item(s) redacted in ${Date.now() - t0} ms.`,
      );
      // ml_status tells the LLM what really happened with the ML layer.
      //   - "disabled":       user has ML off in config; regex-only result
      //   - "ok":             ML ran, names/dates/addresses included
      //   - "error":          ML enabled but unreachable — see ml_error
      const ml_status: "disabled" | "ok" | "error" =
        !useModel ? "disabled" : (modelSpans ? "ok" : "error");

      return {
        anonymized: result.anonymized,
        counts: result.counts,
        ...(args.include_mapping === false ? {} : { mapping: result.mapping }),
        ml_status,
        ...(mlError ? { ml_error: mlError } : {}),
      };
    },
  });

  return [anonymizeText];
}