terminal

import { tool } from "@lmstudio/sdk";
import { z } from "zod";
import { getSystemInfo, clearSystemInfoCache, getNetworkInterfaces } from "../systemUtilities";
import { execCommand, runProcess, normalizeLineEndings } from "../executor";
import { logger } from "../logger";

// Helper for consistent formatting
const formatPct = (n: number) => (n / 100).toFixed(1) + "%";
const formatBytes = (bytes: number) => {
  if (bytes === 0) return "0 Bytes";
  const k = 1024;
  const sizes = ["Bytes", "KB", "MB", "GB", "TB"];
  const i = Math.floor(Math.log(bytes) / Math.log(k));
  return parseFloat((bytes / Math.pow(k, i)).toFixed(2)) + " " + sizes[i];
};

/**
 * Check if exitCode indicates success.
 * On Windows, exitCode can be null for some processes — treat null as 0 (success)
 * unless a signal was received.
 */
function isSuccess(exitCode: number | null, signal?: string): boolean {
  if (signal !== null && signal !== undefined) return false;
  return exitCode === null || exitCode === 0;
}

export async function registerSystemTools(): Promise<any[]> {
  const tools: any[] = [];

  tools.push(tool({
    name: "system_info",
    description: "Get comprehensive system diagnostics.",
    parameters: {},
    implementation: async () => {
      const info = getSystemInfo();
      return {
        success: true,
        output: [
          `Platform: ${info.platform}`,
          `Architecture: ${info.arch}`,
          `CPU Count: ${info.cpuCount}`,
          `Total Memory: ${formatBytes(info.totalMemory)}`,
          `Free Memory: ${formatBytes(info.freeMemory)}`,
          `Used Memory: ${formatPct(info.memoryPercent)}`,
          `Uptime: ${info.uptime.toFixed(0)} seconds`,
          `Hostname: ${info.hostname}`,
          `Load Average: ${info.loadAvg.load1}, ${info.loadAvg.load5}, ${info.loadAvg.load15}`,
          `Temp Dir: ${info.tempDir}`,
          `Home Dir: ${info.homeDir}`,
          "--- Per-Core CPU Usage ---",
          ...info.cpus.map((c, i) => `  Core ${i}: ${c.model} @ ${c.speed} MHz — ${formatPct(c.usage)}`),
        ].join("\n"),
      };
    },
  }));

  tools.push(tool({
    name: "list_processes",
    description: "List running processes.",
    parameters: { limit: z.number().optional().describe("Max processes to return") },
    implementation: async ({ limit }) => {
      // Clamp limit to prevent command injection via large numbers
      const max = Math.min(Math.max(1, Number(limit) || 50), 10000);
      try {
        if (process.platform === "win32") {
          const result = await runProcess("powershell", [
            "-Command",
            `Get-Process | Select-Object -First ${max} | Format-Table Id, ProcessName, CPU, WorkingSet -AutoSize -Wrap`,
          ], { timeoutMs: 10000 });
          // Normalize Windows line endings and trim whitespace
          const output = normalizeLineEndings(result.stdout).trim();
          return { success: isSuccess(result.exitCode, result.signal), output, process_count: max };
        } else {
          const result = await runProcess("ps", ["aux"], { timeoutMs: 10000 });
          const lines = result.stdout.split("\n").slice(0, max + 1); // header + max entries
          return { success: isSuccess(result.exitCode, result.signal), output: lines.join("\n"), process_count: Math.min(max, lines.length - 1) };
        }
      } catch (err: any) {
        logger.error('list_processes failed', { error: err.message });
        return { success: false, output: "", error: err.message };
      }
    },
  }));

  tools.push(tool({
    name: "kill_process",
    description: "Terminate a process by its PID.",
    parameters: { pid: z.number().describe("Process ID"), signal: z.string().optional().describe("Signal") },
    implementation: async ({ pid, signal }) => {
      try {
        // Validate PID is a positive integer to prevent injection
        if (!Number.isInteger(pid) || pid <= 0) {
          return { success: false, pid, signal, error: `Invalid PID: ${pid}. Must be a positive integer.` };
        }

        let result: { exitCode: number | null; stderr: string };
        if (process.platform === "win32") {
          // Use run_process with array args to avoid shell interpretation
          const procResult = await runProcess("taskkill", ["/PID", String(pid), "/F"], {
            timeoutMs: 10000,
          });
          result = { exitCode: procResult.exitCode, stderr: procResult.stderr };
        } else {
          const args = signal ? [`-${signal}`, String(pid)] : [String(pid)];
          const procResult = await runProcess("kill", args, { timeoutMs: 10000 });
          result = { exitCode: procResult.exitCode, stderr: procResult.stderr };
        }
        return {
          success: isSuccess(result.exitCode, undefined),
          pid,
          signal,
          error: !isSuccess(result.exitCode, undefined) ? result.stderr : undefined,
        };
      } catch (err: any) {
        logger.error('kill_process failed', { error: err.message });
        return { success: false, pid, signal, error: err.message };
      }
    },
  }));

  tools.push(tool({
    name: "check_network",
    description: "Test network connectivity.",
    parameters: {
      host: z.string().describe("Host"),
      port: z.number().optional().describe("Port"),
      timeout_ms: z.number().optional().describe("Timeout"),
    },
    implementation: async ({ host, port, timeout_ms }) => {
      const p = port ?? 80;
      const t = timeout_ms ?? 5000;
      try {
        if (process.platform === "win32") {
          // Use PowerShell TCP check — success is indicated by stdout containing "True"
          // (exitCode may be null even on success due to .NET exception handling)
          const safeHost = host.replace(/'/g, "''");
          const script = `[Net.Sockets.TcpClient]::new().ConnectAsync('${safeHost}', ${p}).Wait(${t}); $true`;
          const result = await runProcess("powershell", ["-Command", script], { timeoutMs: t + 2000 });
          const output = normalizeLineEndings(result.stdout).trim();
          // Success if stdout contains "True" (case-insensitive) regardless of exitCode
          const connected = output.toLowerCase().includes("true");
          return {
            success: connected,
            host,
            port: p,
            error: connected ? undefined : (result.stderr || "Connection failed"),
          };
        } else {
          // Sanitize host to prevent injection — only allow alphanumeric, dots, hyphens
          const sanitizedHost = host.replace(/[^a-zA-Z0-9._-]/g, '');
          if (!sanitizedHost) return { success: false, host, port: p, error: "Invalid hostname" };
          const secs = Math.ceil(t / 1000);
          const bashCmd = `timeout ${secs}s bash -c 'echo >/dev/tcp/${sanitizedHost}/${p}' 2>/dev/null`;
          const result = await execCommand(bashCmd, { timeoutMs: t + 2000 });
          return {
            success: isSuccess(result.exitCode, undefined),
            host,
            port: p,
            error: !isSuccess(result.exitCode, undefined) ? result.stderr || result.stdout.trim() : undefined,
          };
        }
      } catch (err: any) {
        logger.error('check_network failed', { error: err.message });
        return { success: false, host, port: p, error: err.message };
      }
    },
  }));

  tools.push(tool({
    name: "get_network_info",
    description: "Get network interfaces.",
    parameters: {},
    implementation: async () => {
      try {
        const interfaces = getNetworkInterfaces();
        const output = interfaces.map(iface =>
          `${iface.name} ${iface.family}: ${iface.address}/${iface.cidr || 'N/A'} MAC:${iface.mac}${iface.internal ? ' (Internal)' : ''}`
        ).join("\n") || "No network interfaces found.";

        return { success: true, output };
      } catch (err: any) {
        logger.error('get_network_info failed', { error: err.message });
        return { success: false, output: "", error: err.message };
      }
    },
  }));

  tools.push(tool({
    name: "get_disk_info",
    description: "Get disk/volume info.",
    parameters: {},
    implementation: async () => {
      try {
        if (process.platform === "win32") {
          const result = await runProcess("powershell", [
            "-Command",
            "Get-CimInstance Win32_LogicalDisk | Select-Object DeviceID, @{N='SizeGB';E={[math]::Round($_.Size/1GB,2)}}, @{N='FreeGB';E={[math]::Round($_.FreeSpace/1GB,2)}} | Format-Table -AutoSize",
          ], { timeoutMs: 10000 });
          const output = normalizeLineEndings(result.stdout).trim();
          return { success: isSuccess(result.exitCode, result.signal), output };
        } else {
          const result = await runProcess("df", ["-h"], { timeoutMs: 10000 });
          return { success: isSuccess(result.exitCode, result.signal), output: result.stdout };
        }
      } catch (err: any) {
        logger.error('get_disk_info failed', { error: err.message });
        return { success: false, output: "", error: err.message };
      }
    },
  }));

  return tools;
}