#!/usr/bin/env node
// MCP stdio server exposing figma_screenshot + figma_compare.
// Spawned per-build by AiControl.Figma's ClaudeCliLauncher — lives only for
// the duration of a single Claude build and is killed when the CLI exits.
//
// Protocol: MCP over newline-delimited JSON-RPC on stdin/stdout.
// Requires: puppeteer-core (uses system Chrome/Edge), pixelmatch, pngjs.
// The app does NOT bundle Chromium — we auto-detect user's Chrome or Edge.

import { readFileSync, writeFileSync, existsSync, mkdirSync } from 'node:fs';
import { dirname, resolve, basename } from 'node:path';
import { tmpdir } from 'node:os';
import readline from 'node:readline';
import { pathToFileURL } from 'node:url';
import { PNG } from 'pngjs';
import pixelmatch from 'pixelmatch';
import puppeteer from 'puppeteer-core';

// ---- Chrome/Edge discovery ---------------------------------------------------

const EXE_CANDIDATES = [
    process.env.PUPPETEER_EXECUTABLE_PATH,
    'C:/Program Files/Google/Chrome/Application/chrome.exe',
    'C:/Program Files (x86)/Google/Chrome/Application/chrome.exe',
    process.env.LOCALAPPDATA + '/Google/Chrome/Application/chrome.exe',
    'C:/Program Files (x86)/Microsoft/Edge/Application/msedge.exe',
    'C:/Program Files/Microsoft/Edge/Application/msedge.exe',
].filter(Boolean);

function findBrowser() {
    for (const p of EXE_CANDIDATES) {
        try { if (existsSync(p)) return p; } catch { }
    }
    return null;
}

let browserPromise = null;
function getBrowser() {
    if (browserPromise) return browserPromise;
    const exe = findBrowser();
    if (!exe) throw new Error('No Chrome or Edge installation found. Install one and retry.');
    browserPromise = puppeteer.launch({
        executablePath: exe,
        headless: 'new',
        args: ['--no-sandbox', '--disable-dev-shm-usage', '--hide-scrollbars'],
    });
    return browserPromise;
}

// ---- Tool implementations ----------------------------------------------------

function ensureDir(p) { mkdirSync(dirname(p), { recursive: true }); }

function readPng(file) {
    const data = readFileSync(file);
    return PNG.sync.read(data);
}

async function screenshot({ html_path, width = 1280, height = 800, full_page = true, out_png }) {
    if (!html_path) throw new Error('html_path is required');
    const absHtml = resolve(html_path);
    if (!existsSync(absHtml)) throw new Error(`html_path not found: ${absHtml}`);
    const url = pathToFileURL(absHtml).toString();

    const outPath = out_png
        ? resolve(out_png)
        : resolve(dirname(absHtml), 'screenshot.png');
    ensureDir(outPath);

    const browser = await getBrowser();
    const page = await browser.newPage();
    try {
        await page.setViewport({ width, height, deviceScaleFactor: 1 });
        await page.goto(url, { waitUntil: 'networkidle2', timeout: 60_000 });
        // Give CDN-loaded react/babel a moment to render after networkidle.
        await new Promise(r => setTimeout(r, 400));
        const buf = await page.screenshot({ type: 'png', fullPage: full_page });
        writeFileSync(outPath, buf);
        const png = PNG.sync.read(buf);
        return { png_path: outPath, width: png.width, height: png.height };
    } finally {
        await page.close().catch(() => { });
    }
}

async function compare({ image_a, image_b, out_diff_png, threshold = 0.1 }) {
    if (!image_a || !image_b) throw new Error('image_a and image_b are required');
    const a = readPng(resolve(image_a));
    const b = readPng(resolve(image_b));

    if (a.width !== b.width || a.height !== b.height) {
        return {
            match_percent: 0,
            pixels_different: a.width * a.height + b.width * b.height,
            total_pixels: a.width * a.height,
            note: `Dimensions differ (A: ${a.width}x${a.height}, B: ${b.width}x${b.height}). Re-render one image at the other's size before comparing.`,
            diff_png_path: null,
        };
    }

    const { width, height } = a;
    const diff = new PNG({ width, height });
    const diffPixels = pixelmatch(a.data, b.data, diff.data, width, height, { threshold, includeAA: false });
    const total = width * height;
    const matchPct = ((total - diffPixels) / total) * 100;

    const outPath = out_diff_png
        ? resolve(out_diff_png)
        : resolve(dirname(image_a), 'diff.png');
    ensureDir(outPath);
    writeFileSync(outPath, PNG.sync.write(diff));

    return {
        match_percent: Number(matchPct.toFixed(3)),
        pixels_different: diffPixels,
        total_pixels: total,
        diff_png_path: outPath,
    };
}

// ---- MCP JSON-RPC wiring -----------------------------------------------------

const TOOLS = [
    {
        name: 'figma_screenshot',
        description: 'Render a local HTML file to a PNG using the system Chrome/Edge browser and save it to disk. Use this to render the generated design so it can be compared against the source Figma frame.',
        inputSchema: {
            type: 'object',
            properties: {
                html_path: { type: 'string', description: 'Absolute path to the HTML file to render.' },
                width: { type: 'number', description: 'Viewport width in CSS pixels. Default 1280.' },
                height: { type: 'number', description: 'Viewport height in CSS pixels. Default 800.' },
                full_page: { type: 'boolean', description: 'Capture the full scrollable page. Default true.' },
                out_png: { type: 'string', description: 'Optional output path. Defaults to screenshot.png next to html_path.' },
            },
            required: ['html_path'],
        },
    },
    {
        name: 'figma_compare',
        description: 'Pixel-diff two PNG images using pixelmatch. Returns match_percent (0-100) plus a diff image showing red-highlighted mismatched pixels. Dimensions must match.',
        inputSchema: {
            type: 'object',
            properties: {
                image_a: { type: 'string', description: 'Absolute path to the rendered screenshot PNG.' },
                image_b: { type: 'string', description: 'Absolute path to the source Figma frame PNG.' },
                out_diff_png: { type: 'string', description: 'Optional output path for the diff image. Defaults to diff.png next to image_a.' },
                threshold: { type: 'number', description: 'Per-pixel color threshold 0..1 (lower = stricter). Default 0.1.' },
            },
            required: ['image_a', 'image_b'],
        },
    },
];

function respond(id, result) {
    process.stdout.write(JSON.stringify({ jsonrpc: '2.0', id, result }) + '\n');
}
function errorReply(id, code, message) {
    process.stdout.write(JSON.stringify({ jsonrpc: '2.0', id, error: { code, message } }) + '\n');
}

async function handleToolCall(name, args) {
    let result;
    if (name === 'figma_screenshot') result = await screenshot(args || {});
    else if (name === 'figma_compare') result = await compare(args || {});
    else throw new Error(`Unknown tool: ${name}`);
    return {
        content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
    };
}

const rl = readline.createInterface({ input: process.stdin, crlfDelay: Infinity });
rl.on('line', async (line) => {
    if (!line.trim()) return;
    let msg;
    try { msg = JSON.parse(line); } catch { return; }
    const { id, method, params } = msg;

    try {
        if (method === 'initialize') {
            respond(id, {
                protocolVersion: '2024-11-05',
                capabilities: { tools: {} },
                serverInfo: { name: 'aicontrol-figma-build', version: '1.0.0' },
            });
        } else if (method === 'notifications/initialized') {
            // no response for notifications
        } else if (method === 'tools/list') {
            respond(id, { tools: TOOLS });
        } else if (method === 'tools/call') {
            const out = await handleToolCall(params?.name, params?.arguments);
            respond(id, out);
        } else if (method === 'ping') {
            respond(id, {});
        } else {
            errorReply(id, -32601, `Method not found: ${method}`);
        }
    } catch (e) {
        // For tools/call, MCP convention is to return isError content, not a JSON-RPC error.
        if (method === 'tools/call') {
            respond(id, {
                content: [{ type: 'text', text: `Error: ${e.message || String(e)}` }],
                isError: true,
            });
        } else {
            errorReply(id ?? null, -32603, e.message || String(e));
        }
    }
});

process.on('SIGINT', async () => {
    try { if (browserPromise) (await browserPromise).close(); } catch { }
    process.exit(0);
});
process.on('SIGTERM', async () => {
    try { if (browserPromise) (await browserPromise).close(); } catch { }
    process.exit(0);
});
