import { z } from 'zod';
import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { config } from '../../config.js';
import { errorResult, jsonResult, toErrorMessage } from '../shared.js';
import { dockerManager } from './dockerManager.js';

export function registerDockerExecTool(server: McpServer): void {
  server.registerTool(
    'docker_exec',
    {
      title: 'Docker Exec',
      description:
        'Run a one-off command inside a running Docker container started by this server and return stdout/stderr.',
      inputSchema: {
        containerId: z.string().min(1).describe('Container ID (short or long) returned by docker_run_container.'),
        command: z.array(z.string()).min(1).describe('Command and arguments to execute inside the container.'),
        workingDir: z.string().optional().describe('Working directory inside the container.'),
        env: z.record(z.string()).optional().describe('Additional environment variables for the exec.'),
      },
    },
    async ({ containerId, command, workingDir, env }) => {
      try {
        await dockerManager.ensureReachable();
        const docker = dockerManager.getDocker();
        const exec = await docker.getContainer(containerId).exec({
          Cmd: command,
          WorkingDir: workingDir,
          Env: env ? Object.entries(env).map(([k, v]) => `${k}=${v}`) : undefined,
          AttachStdout: true,
          AttachStderr: true,
        });
        const stream = await exec.start({ Detach: false, Tty: false });
        const chunks: Buffer[] = [];
        await new Promise<void>((resolve, reject) => {
          stream.on('data', (chunk: Buffer) => chunks.push(chunk));
          stream.on('end', resolve);
          stream.on('error', reject);
        });

        const raw = Buffer.concat(chunks).toString('utf8');
        const result = await exec.inspect();
        return jsonResult({
          containerId,
          command,
          exitCode: result.ExitCode ?? null,
          output: raw.slice(-config.maxToolOutputChars),
        });
      } catch (error) {
        return errorResult(toErrorMessage(error));
      }
    },
  );
}
