import { text, tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { join } from "path";
import { z } from "zod";
import si from 'systeminformation';

export async function toolsProvider(ctl: ToolsProviderController) {
  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 });
    },
  });
  tools.push(getSystemInfoTool);
  return tools;
}

export interface BasicSystemInfo {
  osDistro: string;
  arch: string;
}

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

export type SystemInfoResult = BasicSystemInfo & {
  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 filterData(data: Record<string, any>) {
  const invalidStrings = [
    "unknown",
    "n/a",
    "not available",
    "",
    "system version",
    "system serial number",
    "system product name",
  ];

  return Object.fromEntries(
    Object.entries(data).filter(([_, value]) => {
      if (typeof value === "string") {
        return !invalidStrings.includes(value.toLowerCase().trim());
      }
      if (Array.isArray(value)) {
        return value.length > 0;
      }
      return value !== null && value !== undefined;
    })
  );
}

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];
}