import { z } from 'zod';
import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { writeBaseline, listBaselines } from '../../lib/baselines.js';
import { errorResult, jsonResult, toErrorMessage } from '../shared.js';
import { browserManager } from './browserManager.js';

export function registerBrowserScreenshotBaselineTool(server: McpServer): void {
  server.registerTool(
    'browser_screenshot_baseline',
    {
      title: 'Save Screenshot Baseline',
      description:
        'Capture the current page (or a CSS selector) and save it as a PNG baseline under `.web-dev-mcp/baselines/{name}.png` for later visual diffs.',
      inputSchema: {
        name: z
          .string()
          .min(1)
          .describe('Baseline name (no path separators), e.g. "home" or "login-form".'),
        selector: z.string().optional().describe('Optional CSS selector; if omitted, captures the page.'),
        fullPage: z
          .boolean()
          .optional()
          .default(false)
          .describe('Capture the full scrollable page when selector is omitted.'),
      },
    },
    async ({ name, selector, fullPage }) => {
      try {
        const session = await browserManager.getSession();
        let buffer: Buffer;
        if (selector) {
          const element = await session.page.locator(selector).first();
          await element.waitFor({ state: 'visible', timeout: 5000 });
          buffer = await element.screenshot({ type: 'png' });
        } else {
          buffer = await session.page.screenshot({ fullPage, type: 'png' });
        }

        const saved = writeBaseline(name, buffer);
        return jsonResult({
          name,
          selector: selector ?? null,
          fullPage: selector ? false : fullPage,
          sizeBytes: buffer.length,
          relativePath: saved.relativePath,
          baselines: listBaselines(),
        });
      } catch (error) {
        return errorResult(toErrorMessage(error));
      }
    },
  );
}