import fs from 'node:fs';
import path from 'node:path';
import { z } from 'zod';
import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { config } from '../../config.js';
import { resolveWorkspacePath } from '../../lib/sandbox.js';
import { errorResult, jsonResult, toErrorMessage } from '../shared.js';
import { dockerManager } from './dockerManager.js';

export function registerDockerBuildImageTool(server: McpServer): void {
  server.registerTool(
    'docker_build_image',
    {
      title: 'Docker Build Image',
      description:
        'Build a Docker image from a Dockerfile in the workspace. Streams the build output back and returns the resulting image tag.',
      inputSchema: {
        tag: z.string().min(1).describe('Image tag, e.g. "my-app:latest".'),
        dockerfile: z
          .string()
          .optional()
          .default('Dockerfile')
          .describe('Path to the Dockerfile, relative to the workspace root.'),
        buildContext: z
          .string()
          .optional()
          .default('.')
          .describe('Build context directory, relative to the workspace root.'),
      },
    },
    async ({ tag, dockerfile, buildContext }) => {
      try {
        await dockerManager.ensureReachable();
        const dockerfilePath = resolveWorkspacePath(config.workspaceRoot, dockerfile);
        const contextPath = resolveWorkspacePath(config.workspaceRoot, buildContext);
        if (!fs.existsSync(dockerfilePath)) {
          return errorResult(`Dockerfile not found: ${dockerfilePath}`);
        }

        const stream = await dockerManager.getDocker().buildImage(
          {
            context: contextPath,
            src: [path.basename(dockerfilePath), ...listContextFiles(contextPath)],
          },
          { t: tag, dockerfile: path.basename(dockerfilePath) },
        );

        const output: string[] = [];
        await new Promise<void>((resolve, reject) => {
          dockerManager
            .getDocker()
            .modem.followProgress(
              stream,
              (err, res) => {
                if (err) {
                  reject(err);
                  return;
                }
                const last = res?.[res.length - 1];
                if (last && 'error' in last && last.error) {
                  reject(new Error(String(last.error)));
                  return;
                }
                resolve();
              },
              (event) => {
                const line = typeof event === 'string' ? event : JSON.stringify(event);
                output.push(line);
                if (output.length > config.maxBufferedLogLines) {
                  output.splice(0, output.length - config.maxBufferedLogLines);
                }
              },
            );
        });

        return jsonResult({
          tag,
          dockerfile: dockerfilePath,
          context: contextPath,
          output: output.join('\n').slice(-config.maxToolOutputChars),
        });
      } catch (error) {
        return errorResult(toErrorMessage(error));
      }
    },
  );
}

function listContextFiles(root: string): string[] {
  const files: string[] = [];
  const walk = (dir: string) => {
    const entries = fs.readdirSync(dir, { withFileTypes: true });
    for (const entry of entries) {
      const full = path.join(dir, entry.name);
      const rel = path.relative(root, full);
      if (entry.name === 'node_modules' || entry.name === '.git' || entry.name === 'dist') continue;
      if (entry.isDirectory()) {
        walk(full);
      } else {
        files.push(rel);
      }
    }
  };
  walk(root);
  return files;
}
