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, isAllowListedImage, isLocallyBuiltImage, validateVolumeMount } from './dockerManager.js';

export function registerDockerRunContainerTool(server: McpServer): void {
  server.registerTool(
    'docker_run_container',
    {
      title: 'Docker Run Container',
      description:
        'Run a Docker container from an image with port mappings, env vars, and workspace-confined volume mounts. ' +
        'Locally built images are allowed; non-local images must be on the allow-list unless confirm: true is passed.',
      inputSchema: {
        image: z.string().min(1).describe('Docker image to run, e.g. "my-app:latest" or "nginx".'),
        name: z.string().optional().describe('Optional container name.'),
        ports: z
          .record(z.number().int().positive())
          .optional()
          .describe('Port mapping { hostPort: containerPort }.'),
        env: z.record(z.string()).optional().describe('Environment variables.'),
        volumes: z
          .array(z.string())
          .optional()
          .describe('Volume mounts in "source:target[:mode]" format. Source paths are confined to WORKSPACE_ROOT.'),
        command: z.array(z.string()).optional().describe('Override the container command.'),
        memory: z.string().optional().describe('Memory limit, e.g. "512m". Defaults to 512m.'),
        cpus: z.string().optional().describe('CPU limit, e.g. "1.0". Defaults to 1.0.'),
        confirm: z
          .boolean()
          .optional()
          .describe('Required as true to run an image not on the allow-list.'),
      },
    },
    async ({ image, name, ports, env, volumes, command, memory, cpus, confirm }) => {
      try {
        await dockerManager.ensureReachable();

        if (!isLocallyBuiltImage(image) && !isAllowListedImage(image) && confirm !== true) {
          return errorResult(
            `Image "${image}" is not locally built and not on the allow-list (${config.dockerAllowedBaseImages.join(', ')}). ` +
              'Pass confirm: true to run it anyway.',
          );
        }

        const portBindings: Record<string, Array<{ HostPort: string }>> = {};
        const exposedPorts: Record<string, {}> = {};
        const portMap: Record<number, number> = {};
        if (ports) {
          for (const [hostPort, containerPort] of Object.entries(ports)) {
            const key = `${containerPort}/tcp`;
            portBindings[key] = [{ HostPort: hostPort }];
            exposedPorts[key] = {};
            portMap[Number(hostPort)] = containerPort;
          }
        }

        const binds: string[] = [];
        if (volumes) {
          for (const spec of volumes) {
            const mount = validateVolumeMount(spec);
            binds.push(`${mount.source}:${mount.target}:${mount.mode}`);
          }
        }

        const createOptions = {
          Image: image,
          name,
          Cmd: command,
          Env: env ? Object.entries(env).map(([k, v]) => `${k}=${v}`) : undefined,
          ExposedPorts: Object.keys(exposedPorts).length > 0 ? exposedPorts : undefined,
          HostConfig: {
            PortBindings: Object.keys(portBindings).length > 0 ? portBindings : undefined,
            Binds: binds.length > 0 ? binds : undefined,
            Memory: parseMemory(memory ?? config.defaultContainerMemory),
            NanoCpus: parseCpus(cpus ?? config.defaultContainerCpus),
            // Security guardrails:
            Privileged: false,
            NetworkMode: 'bridge',
            PidMode: undefined,
            IpcMode: undefined,
          },
        };

        const docker = dockerManager.getDocker();
        const container = await docker.createContainer(createOptions);
        await container.start();
        const info = await container.inspect();
        const managed = dockerManager.track(
          container.id,
          image,
          info.Name?.replace(/^\//, '') ?? container.id,
          portMap,
        );
        return jsonResult({ ...managed, shortId: container.id.substring(0, 12) });
      } catch (error) {
        return errorResult(toErrorMessage(error));
      }
    },
  );
}

function parseMemory(value: string): number {
  const match = value.match(/^([0-9.]+)([kmgt]?b?)$/i);
  if (!match || !match[1]) return 512 * 1024 * 1024;
  const num = Number.parseFloat(match[1]);
  const unit = match[2]?.toLowerCase() ?? '';
  switch (unit) {
    case 'k':
    case 'kb':
      return num * 1024;
    case 'm':
    case 'mb':
      return num * 1024 * 1024;
    case 'g':
    case 'gb':
      return num * 1024 * 1024 * 1024;
    case 't':
    case 'tb':
      return num * 1024 * 1024 * 1024 * 1024;
    default:
      return num;
  }
}

function parseCpus(value: string): number {
  const num = Number.parseFloat(value);
  return Number.isNaN(num) ? 1_000_000_000 : Math.round(num * 1_000_000_000);
}
