import { Type } from '@sinclair/typebox';
import type { ToolDefinition } from '@mariozechner/pi-coding-agent';
import { spawn } from 'node:child_process';
import { guardPath, SandboxEscapeError } from './path-guard.js';

const DEFAULT_TIMEOUT_MS = 30_000;
const MAX_OUTPUT_BYTES = 100 * 1024; // 100 KB

/**
 * Commands that are outright blocked for safety.
 * This is a denylist; the agent should be instructed to use
 * the least-privilege command necessary.
 */
const BLOCKED_COMMANDS = new Set([
  'rm',
  'rmdir',
  'mkfs',
  'dd',
  'format',
  'fdisk',
  'parted',
  'shred',
  'wipefs',
  'sudo',
  'su',
  'chown',
  'chmod',
  'passwd',
  'useradd',
  'userdel',
  'groupadd',
  'shutdown',
  'reboot',
  'halt',
  'poweroff',
  'kill',
  'killall',
  'pkill',
  'curl',
  'wget',
  'nc',
  'netcat',
  'ncat',
  'ssh',
  'scp',
  'sftp',
  'rsync',
  'iptables',
  'ip6tables',
  'nft',
  'ufw',
  'firewall-cmd',
  'docker',
  'podman',
  'kubectl',
  'helm',
  'terraform',
  'ansible',
  'crontab',
  'at',
  'batch',
]);

function extractBaseCommand(command: string): string {
  // Extract the first word (the binary name), stripping path
  const trimmed = command.trim();
  const firstToken = trimmed.split(/\s+/)[0] ?? '';
  return firstToken.split('/').pop() ?? firstToken;
}

function runCommand(
  command: string,
  options: { timeoutMs: number; cwd?: string },
): Promise<{ stdout: string; stderr: string; exitCode: number | null; timedOut: boolean }> {
  return new Promise((resolve) => {
    const child = spawn('sh', ['-c', command], {
      cwd: options.cwd,
      stdio: ['ignore', 'pipe', 'pipe'],
      detached: false,
    });

    let stdout = '';
    let stderr = '';
    let timedOut = false;
    let totalBytes = 0;
    let truncated = false;

    child.stdout?.on('data', (chunk: Buffer) => {
      if (truncated) return;
      totalBytes += chunk.length;
      if (totalBytes > MAX_OUTPUT_BYTES) {
        stdout += chunk.subarray(0, MAX_OUTPUT_BYTES - (totalBytes - chunk.length)).toString();
        stdout += '\n[output truncated at 100 KB limit]';
        truncated = true;
        child.kill('SIGTERM');
      } else {
        stdout += chunk.toString();
      }
    });

    child.stderr?.on('data', (chunk: Buffer) => {
      if (stderr.length < MAX_OUTPUT_BYTES) {
        stderr += chunk.toString();
      }
    });

    const timer = setTimeout(() => {
      timedOut = true;
      child.kill('SIGTERM');
      setTimeout(() => {
        try {
          child.kill('SIGKILL');
        } catch {
          // already exited
        }
      }, 2000);
    }, options.timeoutMs);

    child.on('close', (exitCode) => {
      clearTimeout(timer);
      resolve({ stdout, stderr, exitCode, timedOut });
    });

    child.on('error', (err) => {
      clearTimeout(timer);
      resolve({ stdout, stderr: stderr + String(err), exitCode: null, timedOut: false });
    });
  });
}

export function createShellTools(sandboxDir?: string): ToolDefinition[] {
  const defaultCwd = sandboxDir ?? process.cwd();

  const shellExec: ToolDefinition = {
    name: 'shell_exec',
    label: 'Shell Execute',
    description:
      'Execute a shell command with timeout and output limits. Dangerous commands (rm, sudo, docker, etc.) are blocked. Working directory is restricted to the session sandbox.',
    parameters: Type.Object({
      command: Type.String({ description: 'Shell command to execute' }),
      cwd: Type.Optional(
        Type.String({
          description:
            'Working directory for the command (relative to sandbox or absolute within it).',
        }),
      ),
      timeout: Type.Optional(
        Type.Number({ description: 'Timeout in milliseconds (default 30000, max 60000)' }),
      ),
    }),
    async execute(_toolCallId, params) {
      const { command, cwd, timeout } = params as {
        command: string;
        cwd?: string;
        timeout?: number;
      };

      const base = extractBaseCommand(command);
      if (BLOCKED_COMMANDS.has(base)) {
        return {
          content: [
            {
              type: 'text' as const,
              text: `Error: command "${base}" is blocked for safety reasons.`,
            },
          ],
          details: undefined,
        };
      }

      const timeoutMs = Math.min(timeout ?? DEFAULT_TIMEOUT_MS, 60_000);
      let safeCwd: string;
      try {
        safeCwd = guardPath(cwd ?? '.', defaultCwd);
      } catch (err) {
        if (err instanceof SandboxEscapeError) {
          return {
            content: [{ type: 'text' as const, text: `Error: ${err.message}` }],
            details: undefined,
          };
        }
        return {
          content: [{ type: 'text' as const, text: `Error: ${String(err)}` }],
          details: undefined,
        };
      }

      const result = await runCommand(command, {
        timeoutMs,
        cwd: safeCwd,
      });

      if (result.timedOut) {
        return {
          content: [
            {
              type: 'text' as const,
              text: `Command timed out after ${timeoutMs}ms.\nPartial stdout:\n${result.stdout}\nPartial stderr:\n${result.stderr}`,
            },
          ],
          details: undefined,
        };
      }

      const parts: string[] = [];
      if (result.stdout) parts.push(`stdout:\n${result.stdout}`);
      if (result.stderr) parts.push(`stderr:\n${result.stderr}`);
      parts.push(`exit code: ${result.exitCode ?? 'null'}`);

      return {
        content: [{ type: 'text' as const, text: parts.join('\n') }],
        details: undefined,
      };
    },
  };

  return [shellExec];
}