sysinfo

import { text, tool, type Tool } from "@lmstudio/sdk";
import { z } from "zod/v3";
import si from "systeminformation";

export async function toolsProvider() {
  const tools: Tool[] = [];

  const getSystemInfoTool = tool({
    name: "get_system_info",
    description: text`Get information about the user system such as operating system, architecture, cpu, graphics, memory, and battery.`,
    parameters: {
      cpu: z.boolean().optional(),
      graphics: z.boolean().optional(),
      mem: z.boolean().optional(),
      battery: z.boolean().optional(),
    },
    implementation: async ({ cpu, graphics, mem, battery }) => {
      return await getSystemInfo({ cpu, graphics, mem, battery });
    },
  });

  const getCurrentTimeTool = tool({
    name: "get_current_time",
    description: text`Get the current time, date and timezone.`,
    parameters: {},
    implementation: async () => {
      return {
        time: new Date().toLocaleTimeString(),
        date: new Date().toLocaleDateString(),
        timezone: Intl.DateTimeFormat().resolvedOptions().timeZone,
      };
    },
  });
  tools.push(getSystemInfoTool, getCurrentTimeTool);
  return tools;
}

export interface SystemInfoOptions {
  cpu?: boolean;
  graphics?: boolean;
  mem?: boolean;
  battery?: boolean;
}

export type SystemInfoResult = {
  cpu?: object;
  graphics?: object;
  mem?: object;
  battery?: si.Systeminformation.BatteryData | string;
};

export async function getSystemInfo(
  options: SystemInfoOptions = {},
): Promise<SystemInfoResult> {
  const osInfo = await si.osInfo();

  const result: SystemInfoResult = {
    osDistro: osInfo.distro,
    arch: osInfo.arch,
  };

  const promises: Promise<void>[] = [];

  if (options.cpu) {
    promises.push(
      si.cpu().then((cpu) => {
        result.cpu = {
          brand: cpu.manufacturer + " " + cpu.brand,
          speed:
            cpu.speedMin == cpu.speedMax
              ? cpu.speed
              : {
                  current: cpu.speed,
                  min: cpu.speedMin,
                  max: cpu.speedMax,
                },
          cores: cpu.cores,
          physicalCores: cpu.physicalCores,
          performanceCores: cpu.performanceCores,
          efficiencyCores: cpu.efficiencyCores,
          socket: cpu.socket,
          flags: cpu.flags,
          virtualization: cpu.virtualization,
          cache: cpu.cache,
        };
      }),
    );
  }

  if (options.graphics) {
    promises.push(
      si.graphics().then((graphics) => {
        result.graphics = graphics.controllers.map((controller) => ({
          model: controller.model,
          vram: formatSize(controller.vram ?? 0, "mb"),
        }));
      }),
    );
  }

  if (options.mem) {
    promises.push(
      si.mem().then((mem) => {
        result.mem = {
          total: formatSize(mem.total),
          free: formatSize(mem.free),
          used: formatSize(mem.used),
          swapTotal: formatSize(mem.swaptotal),
          swapFree: formatSize(mem.swapfree),
          swapUsed: formatSize(mem.swapused),
        };
      }),
    );
  }

  if (options.battery) {
    promises.push(
      si.battery().then((bat) => {
        if (bat.hasBattery) {
          result.battery = bat;
        } else {
          result.battery = "No battery found";
        }
      }),
    );
  }

  await Promise.all(promises);
  return result;
}

function formatSize(size: number, format_by: string = "b") {
  let bytes = size;
  switch (format_by.toLowerCase()) {
    case "kb":
      bytes = size * 1024;
      break;
    case "mb":
      bytes = size * 1024 * 1024;
      break;
    case "gb":
      bytes = size * 1024 * 1024 * 1024;
      break;
    case "tb":
      bytes = size * 1024 * 1024 * 1024 * 1024;
      break;
    case "b":
    default:
      bytes = size;
      break;
  }

  const units = ["B", "KB", "MB", "GB", "TB"];
  let index = 0;
  while (bytes >= 1024 && index < units.length - 1) {
    bytes /= 1024;
    index++;
  }
  return bytes.toFixed(2) + " " + units[index];
}