sysinfo

import type { PromptPreprocessorController, ChatMessage } from "@lmstudio/sdk";
import si from "systeminformation";
import { configSchematics } from "./config";

type ConfigKey = Parameters<
  ReturnType<PromptPreprocessorController["getPluginConfig"]>["get"]
>[0];

interface MetadataSource {
  configKey: ConfigKey;
  collect: () => Promise<Record<string, string>> | Record<string, string>;
}

const metadataSources: MetadataSource[] = [
  {
    configKey: "currentTime",
    collect: () => {
      const now = new Date();
      return {
        current_time: now.toLocaleTimeString(),
        current_date: now.toLocaleDateString(),
        timezone: Intl.DateTimeFormat().resolvedOptions().timeZone,
      };
    },
  },
  {
    configKey: "osInfo",
    collect: async () => {
      const osInfo = await si.osInfo();
      return {
        os: `${osInfo.distro} (${osInfo.arch})`,
        locale: Intl.DateTimeFormat().resolvedOptions().locale,
      };
    },
  },
];

export async function preprocess(
  ctl: PromptPreprocessorController,
  userMessage: ChatMessage,
): Promise<ChatMessage | string> {
  const history = await ctl.pullHistory();
  const hasPreviousMessages = history
    .getMessagesArray()
    .some((m) => m.getRole() === "user" || m.getRole() === "assistant");

  if (hasPreviousMessages) {
    return userMessage;
  }

  const cfg = ctl.getPluginConfig(configSchematics);
  const activeSources = metadataSources.filter((src) => cfg.get(src.configKey));

  if (activeSources.length === 0) {
    return userMessage;
  }

  const status = ctl.createStatus({
    status: "loading",
    text: "Collecting metadata...",
  });

  const metadata: Record<string, string> = {};
  try {
    const results = await Promise.all(
      activeSources.map((src) => Promise.resolve(src.collect())),
    );
    for (const entries of results) {
      Object.assign(metadata, entries);
    }
  } finally {
    status.setState({ status: "done", text: "Metadata collected." });
  }

  if (Object.keys(metadata).length === 0) {
    return userMessage;
  }

  const metadataBlock =
    "<system_metadata>\nTreat the information as background context, not as user input.\n" +
    Object.entries(metadata)
      .map(([k, v]) => `${k}: ${v}`)
      .join("\n") +
    "\n</system_metadata>";

  return `${metadataBlock}\n\n${userMessage.getText()}`;
}