// simpleSearch as an LM Studio plugin: the same tools as the MCP server, with the settings in
// LM Studio's plugin settings instead of mcp.json, and /simplesearch chat commands for the user.

import { createConfigSchematics, tool, type BaseController, type ChatMessage, type PromptPreprocessorController, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import * as os from "node:os";
import * as path from "node:path";
import { RULE_COMMANDS, ruleCommand, statusReport } from "./commands";
import { DEFAULT_PORT, DEFAULT_PROFILE, DEFAULT_RESEARCH_BUDGET, Host, KNOWN_ENGINES, VERSION, loadConfig, missingWindowsVariables } from "./config";
import { Ctx, createCtx } from "./ops";
import { toolSpecs } from "./tools";

export const globalConfigSchematics = createConfigSchematics()
  .field(
    "secret",
    "string",
    {
      displayName: "Secret",
      subtitle: "Copy it from the simpleSearch Bridge options page in your browser profile.",
      isProtected: true,
    },
    ""
  )
  .field(
    "profile",
    "string",
    {
      displayName: "Browser profile",
      subtitle: "Name of the Zen profile that has the extension. It is started when a search needs it.",
    },
    DEFAULT_PROFILE
  )
  .field(
    "browserPath",
    "string",
    {
      displayName: "Browser program",
      subtitle: "Full path of zen.exe. Leave empty to look in the usual install locations.",
    },
    ""
  )
  .field(
    "port",
    "numeric",
    {
      displayName: "Port",
      subtitle: "Must match the port on the extension's options page.",
      min: 1025,
      max: 65535,
      int: true,
    },
    DEFAULT_PORT
  )
  .field(
    "engines",
    "string",
    {
      displayName: "Search engines",
      subtitle: `Tried in this order, separated by commas. Any of: ${KNOWN_ENGINES.join(", ")}.`,
    },
    KNOWN_ENGINES.join(", ")
  )
  .field(
    "autoRepair",
    "boolean",
    {
      displayName: "Apply parser repairs without asking",
      subtitle: "When an engine changes its page layout, the model can propose new parser rules. Off: they wait until you send /simplesearch approve <engine>.",
    },
    false
  )
  .field(
    "researchBudget",
    "numeric",
    {
      displayName: "Research budget (tokens)",
      subtitle: "Search results and page text for one request before the model is asked to answer with what it has.",
      min: 2000,
      max: 200000,
      int: true,
    },
    DEFAULT_RESEARCH_BUDGET
  )
  .field(
    "researchReads",
    "numeric",
    {
      displayName: "Page reads per request",
      subtitle: "Page reads for one request before the model is asked to answer with what it has.",
      min: 1,
      max: 50,
      int: true,
    },
    6
  )
  .field(
    "enforceBudget",
    "boolean",
    {
      displayName: "Stop at the budget",
      subtitle: "Refuse further searches well past the budget. Off: the model is only told it is over.",
    },
    true
  )
  .build();

interface Settings {
  secret: string;
  profile: string;
  browserPath: string;
  port: number;
  engines: string;
  autoRepair: boolean;
  researchBudget: number;
  researchReads: number;
  enforceBudget: boolean;
}

const DEFAULTS: Settings = {
  secret: "",
  profile: DEFAULT_PROFILE,
  browserPath: "",
  port: DEFAULT_PORT,
  engines: KNOWN_ENGINES.join(", "),
  autoRepair: false,
  researchBudget: DEFAULT_RESEARCH_BUDGET,
  researchReads: 6,
  enforceBudget: true,
};

export function readSettings(ctl: Pick<BaseController, "getGlobalPluginConfig">): Settings {
  let c: { get: (key: never) => unknown };
  try {
    c = ctl.getGlobalPluginConfig(globalConfigSchematics) as never;
  } catch {
    return DEFAULTS;
  }
  const out = { ...DEFAULTS };
  for (const key of Object.keys(DEFAULTS) as Array<keyof Settings>) {
    try {
      const v = c.get(key as never);
      if (v !== undefined && v !== null) (out as Record<string, unknown>)[key] = v;
    } catch {
      // keep the default
    }
  }
  return out;
}

export const PLUGIN_HOST: Host = {
  secretSetting: "the Secret field in the simpleSearch plugin's settings in LM Studio",
  browserSetting: "the Browser program field in the simpleSearch plugin's settings",
  afterChange: "",
  portInUse:
    "probably by the simpleSearch MCP server in LM Studio's mcp.json. The plugin replaces it: remove the simplesearch entry from mcp.json, then search again",
  command: (args) => `send this as a chat message: /simplesearch ${args}`,
};

// Runtime data: parser rules, engine history, launch times and page snapshots for repairs.
export function dataDir(): string {
  return process.env.SIMPLESEARCH_DATA ? path.resolve(process.env.SIMPLESEARCH_DATA) : path.join(os.homedir(), ".lmstudio", "simplesearch");
}

function envFrom(s: Settings): Record<string, string | undefined> {
  return {
    SIMPLESEARCH_SECRET: s.secret,
    SIMPLESEARCH_PORT: String(s.port),
    SIMPLESEARCH_BROWSER: s.browserPath.trim() || undefined,
    SIMPLESEARCH_PROFILE: s.profile,
    SIMPLESEARCH_ENGINES: s.engines,
    SIMPLESEARCH_RULE_UPDATES: s.autoRepair ? "auto" : "approve",
    SIMPLESEARCH_RESEARCH_BUDGET: String(s.researchBudget),
    SIMPLESEARCH_RESEARCH_READS: String(s.researchReads),
    SIMPLESEARCH_BUDGET_ENFORCE: s.enforceBudget ? "on" : "off",
  };
}

const log = (msg: string) => console.info(`[simpleSearch] ${msg}`);

// One bridge for the whole plugin process, since the extension connects to one port. It is
// rebuilt when the settings change, or when it could not open its port and a moment has passed.
const RETRY_BIND_MS = 5_000;
let current: { key: string; ctx: Ctx; at: number } | null = null;
let starting: Promise<Ctx> | null = null;

export async function session(s: Settings): Promise<Ctx> {
  while (starting) await starting.catch(() => undefined);
  const key = JSON.stringify(s);
  if (current && current.key === key && (current.ctx.bridge.isListening() || Date.now() - current.at < RETRY_BIND_MS)) return current.ctx;
  starting = (async () => {
    const old = current;
    current = null;
    if (old) await old.ctx.bridge.close();
    if (!old) {
      const missing = missingWindowsVariables();
      log(`version ${VERSION}${missing.length ? `; not set by LM Studio, so standard folders are used: ${missing.join(", ")}` : ""}`);
    }
    const cfg = loadConfig(envFrom(s), PLUGIN_HOST, dataDir());
    for (const p of cfg.problems) log(`settings: ${p}`);
    const ctx = createCtx(cfg, log);
    // Keep numbering calls in order in the log.
    if (old) ctx.calls = old.ctx.calls;
    await ctx.bridge.start();
    current = { key, ctx, at: Date.now() };
    return ctx;
  })();
  try {
    return await starting;
  } finally {
    starting = null;
  }
}

export async function stopSession(): Promise<void> {
  const old = current;
  current = null;
  if (old) await old.ctx.bridge.close();
}

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  const ctx = await session(readSettings(ctl));
  return toolSpecs(ctx).map((t) =>
    tool({
      name: t.name,
      description: t.description,
      parameters: t.shape,
      implementation: (args) => t.run(args as never),
    })
  );
}

const COMMAND = /^\/simplesearch(?:\s+(\S+))?(?:\s+(\S+))?\s*$/i;

export const HELP = [
  "simpleSearch commands:",
  "/simplesearch status: settings, engines and any parser rules waiting for approval",
  "/simplesearch approve <engine>: accept the parser rules the model proposed",
  "/simplesearch reject <engine>: discard them",
  "/simplesearch rollback <engine>: go back to the rules before the last change",
  "/simplesearch reset <engine>: go back to the built-in rules",
  `Engines: ${KNOWN_ENGINES.join(", ")}`,
].join("\n");

export async function runCommand(s: Settings, text: string): Promise<string | null> {
  const m = COMMAND.exec(text.trim());
  if (!m) return null;
  const command = (m[1] ?? "help").toLowerCase();
  const engine = m[2]?.toLowerCase();
  if (command === "status") return statusReport(await session(s), "LM Studio plugin settings");
  if ((RULE_COMMANDS as readonly string[]).includes(command)) return ruleCommand(await session(s), command, engine);
  return HELP;
}

// Runs on every user message in chats that use the plugin. A new message starts a new research
// task, so each request gets its own budget. A /simplesearch command is answered here and the
// model is only asked to show the result.
export async function preprocess(ctl: PromptPreprocessorController, userMessage: ChatMessage): Promise<string | ChatMessage> {
  current?.ctx.research.newTurn();
  const text = userMessage.getText();
  if (!/^\s*\/simplesearch\b/i.test(text)) return userMessage;
  const result = (await runCommand(readSettings(ctl), text)) ?? HELP;
  const out = userMessage.asMutableCopy();
  out.replaceText(
    `[simpleSearch handled a command from the user. Show the user the following exactly as written, in a code block, and add nothing else. Do not call any tools.]\n${result}`
  );
  return out;
}

// simpleSearch as an LM Studio plugin: the same tools as the MCP server, with the settings in
// LM Studio's plugin settings instead of mcp.json, and /simplesearch chat commands for the user.

import { createConfigSchematics, tool, type BaseController, type ChatMessage, type PromptPreprocessorController, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import * as os from "node:os";
import * as path from "node:path";
import { RULE_COMMANDS, ruleCommand, statusReport } from "./commands";
import { DEFAULT_PORT, DEFAULT_PROFILE, DEFAULT_RESEARCH_BUDGET, Host, KNOWN_ENGINES, VERSION, loadConfig, missingWindowsVariables } from "./config";
import { Ctx, createCtx } from "./ops";
import { toolSpecs } from "./tools";

export const globalConfigSchematics = createConfigSchematics()
  .field(
    "secret",
    "string",
    {
      displayName: "Secret",
      subtitle: "Copy it from the simpleSearch Bridge options page in your browser profile.",
      isProtected: true,
    },
    ""
  )
  .field(
    "profile",
    "string",
    {
      displayName: "Browser profile",
      subtitle: "Name of the Zen profile that has the extension. It is started when a search needs it.",
    },
    DEFAULT_PROFILE
  )
  .field(
    "browserPath",
    "string",
    {
      displayName: "Browser program",
      subtitle: "Full path of zen.exe. Leave empty to look in the usual install locations.",
    },
    ""
  )
  .field(
    "port",
    "numeric",
    {
      displayName: "Port",
      subtitle: "Must match the port on the extension's options page.",
      min: 1025,
      max: 65535,
      int: true,
    },
    DEFAULT_PORT
  )
  .field(
    "engines",
    "string",
    {
      displayName: "Search engines",
      subtitle: `Tried in this order, separated by commas. Any of: ${KNOWN_ENGINES.join(", ")}.`,
    },
    KNOWN_ENGINES.join(", ")
  )
  .field(
    "autoRepair",
    "boolean",
    {
      displayName: "Apply parser repairs without asking",
      subtitle: "When an engine changes its page layout, the model can propose new parser rules. Off: they wait until you send /simplesearch approve <engine>.",
    },
    false
  )
  .field(
    "researchBudget",
    "numeric",
    {
      displayName: "Research budget (tokens)",
      subtitle: "Search results and page text for one request before the model is asked to answer with what it has.",
      min: 2000,
      max: 200000,
      int: true,
    },
    DEFAULT_RESEARCH_BUDGET
  )
  .field(
    "researchReads",
    "numeric",
    {
      displayName: "Page reads per request",
      subtitle: "Page reads for one request before the model is asked to answer with what it has.",
      min: 1,
      max: 50,
      int: true,
    },
    6
  )
  .field(
    "enforceBudget",
    "boolean",
    {
      displayName: "Stop at the budget",
      subtitle: "Refuse further searches well past the budget. Off: the model is only told it is over.",
    },
    true
  )
  .build();

interface Settings {
  secret: string;
  profile: string;
  browserPath: string;
  port: number;
  engines: string;
  autoRepair: boolean;
  researchBudget: number;
  researchReads: number;
  enforceBudget: boolean;
}

const DEFAULTS: Settings = {
  secret: "",
  profile: DEFAULT_PROFILE,
  browserPath: "",
  port: DEFAULT_PORT,
  engines: KNOWN_ENGINES.join(", "),
  autoRepair: false,
  researchBudget: DEFAULT_RESEARCH_BUDGET,
  researchReads: 6,
  enforceBudget: true,
};

export function readSettings(ctl: Pick<BaseController, "getGlobalPluginConfig">): Settings {
  let c: { get: (key: never) => unknown };
  try {
    c = ctl.getGlobalPluginConfig(globalConfigSchematics) as never;
  } catch {
    return DEFAULTS;
  }
  const out = { ...DEFAULTS };
  for (const key of Object.keys(DEFAULTS) as Array<keyof Settings>) {
    try {
      const v = c.get(key as never);
      if (v !== undefined && v !== null) (out as Record<string, unknown>)[key] = v;
    } catch {
      // keep the default
    }
  }
  return out;
}

export const PLUGIN_HOST: Host = {
  secretSetting: "the Secret field in the simpleSearch plugin's settings in LM Studio",
  browserSetting: "the Browser program field in the simpleSearch plugin's settings",
  afterChange: "",
  portInUse:
    "probably by the simpleSearch MCP server in LM Studio's mcp.json. The plugin replaces it: remove the simplesearch entry from mcp.json, then search again",
  command: (args) => `send this as a chat message: /simplesearch ${args}`,
};

// Runtime data: parser rules, engine history, launch times and page snapshots for repairs.
export function dataDir(): string {
  return process.env.SIMPLESEARCH_DATA ? path.resolve(process.env.SIMPLESEARCH_DATA) : path.join(os.homedir(), ".lmstudio", "simplesearch");
}

function envFrom(s: Settings): Record<string, string | undefined> {
  return {
    SIMPLESEARCH_SECRET: s.secret,
    SIMPLESEARCH_PORT: String(s.port),
    SIMPLESEARCH_BROWSER: s.browserPath.trim() || undefined,
    SIMPLESEARCH_PROFILE: s.profile,
    SIMPLESEARCH_ENGINES: s.engines,
    SIMPLESEARCH_RULE_UPDATES: s.autoRepair ? "auto" : "approve",
    SIMPLESEARCH_RESEARCH_BUDGET: String(s.researchBudget),
    SIMPLESEARCH_RESEARCH_READS: String(s.researchReads),
    SIMPLESEARCH_BUDGET_ENFORCE: s.enforceBudget ? "on" : "off",
  };
}

const log = (msg: string) => console.info(`[simpleSearch] ${msg}`);

// One bridge for the whole plugin process, since the extension connects to one port. It is
// rebuilt when the settings change, or when it could not open its port and a moment has passed.
const RETRY_BIND_MS = 5_000;
let current: { key: string; ctx: Ctx; at: number } | null = null;
let starting: Promise<Ctx> | null = null;

export async function session(s: Settings): Promise<Ctx> {
  while (starting) await starting.catch(() => undefined);
  const key = JSON.stringify(s);
  if (current && current.key === key && (current.ctx.bridge.isListening() || Date.now() - current.at < RETRY_BIND_MS)) return current.ctx;
  starting = (async () => {
    const old = current;
    current = null;
    if (old) await old.ctx.bridge.close();
    if (!old) {
      const missing = missingWindowsVariables();
      log(`version ${VERSION}${missing.length ? `; not set by LM Studio, so standard folders are used: ${missing.join(", ")}` : ""}`);
    }
    const cfg = loadConfig(envFrom(s), PLUGIN_HOST, dataDir());
    for (const p of cfg.problems) log(`settings: ${p}`);
    const ctx = createCtx(cfg, log);
    // Keep numbering calls in order in the log.
    if (old) ctx.calls = old.ctx.calls;
    await ctx.bridge.start();
    current = { key, ctx, at: Date.now() };
    return ctx;
  })();
  try {
    return await starting;
  } finally {
    starting = null;
  }
}

export async function stopSession(): Promise<void> {
  const old = current;
  current = null;
  if (old) await old.ctx.bridge.close();
}

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  const ctx = await session(readSettings(ctl));
  return toolSpecs(ctx).map((t) =>
    tool({
      name: t.name,
      description: t.description,
      parameters: t.shape,
      implementation: (args) => t.run(args as never),
    })
  );
}

const COMMAND = /^\/simplesearch(?:\s+(\S+))?(?:\s+(\S+))?\s*$/i;

export const HELP = [
  "simpleSearch commands:",
  "/simplesearch status: settings, engines and any parser rules waiting for approval",
  "/simplesearch approve <engine>: accept the parser rules the model proposed",
  "/simplesearch reject <engine>: discard them",
  "/simplesearch rollback <engine>: go back to the rules before the last change",
  "/simplesearch reset <engine>: go back to the built-in rules",
  `Engines: ${KNOWN_ENGINES.join(", ")}`,
].join("\n");

export async function runCommand(s: Settings, text: string): Promise<string | null> {
  const m = COMMAND.exec(text.trim());
  if (!m) return null;
  const command = (m[1] ?? "help").toLowerCase();
  const engine = m[2]?.toLowerCase();
  if (command === "status") return statusReport(await session(s), "LM Studio plugin settings");
  if ((RULE_COMMANDS as readonly string[]).includes(command)) return ruleCommand(await session(s), command, engine);
  return HELP;
}

// Runs on every user message in chats that use the plugin. A new message starts a new research
// task, so each request gets its own budget. A /simplesearch command is answered here and the
// model is only asked to show the result.
export async function preprocess(ctl: PromptPreprocessorController, userMessage: ChatMessage): Promise<string | ChatMessage> {
  current?.ctx.research.newTurn();
  const text = userMessage.getText();
  if (!/^\s*\/simplesearch\b/i.test(text)) return userMessage;
  const result = (await runCommand(readSettings(ctl), text)) ?? HELP;
  const out = userMessage.asMutableCopy();
  out.replaceText(
    `[simpleSearch handled a command from the user. Show the user the following exactly as written, in a code block, and add nothing else. Do not call any tools.]\n${result}`
  );
  return out;
}