vite.ui.config.ts

import { defineConfig, Plugin, PluginOption } from "vite";
import react from "@vitejs/plugin-react";
import { viteSingleFile } from "vite-plugin-singlefile";
import { nodePolyfills } from "vite-plugin-node-polyfills";
import { readFileSync, writeFileSync, mkdirSync, existsSync, readdirSync, statSync, rmSync } from "fs";
import { join, resolve } from "path";

const componentsDir = resolve(__dirname, "src/ui/components");
// Use node_modules/.cache for generated HTML entries - these are build artifacts, not source files
const entriesDir = resolve(__dirname, "node_modules/.cache/mongodb-mcp-server/ui-entries");
const templatePath = resolve(__dirname, "src/ui/build/template.html");
const mountPath = resolve(__dirname, "src/ui/build/mount.tsx");
const generatedDir = resolve(__dirname, "src/ui/lib");
const uiDistPath = resolve(__dirname, "dist/ui");

// Converts PascalCase to kebab-case: "ListDatabases" -> "list-databases"
function toKebabCase(pascalCase: string): string {
    return pascalCase
        .replace(/([a-z0-9])([A-Z])/g, "$1-$2")
        .replace(/([A-Z])([A-Z][a-z])/g, "$1-$2")
        .toLowerCase();
}

// Discovers component directories and builds tool name mappings
function discoverComponents(): { components: string[]; toolToComponentMap: Record<string, string> } {
    const components: string[] = [];
    const toolToComponentMap: Record<string, string> = {};

    for (const entry of readdirSync(componentsDir)) {
        const entryPath = join(componentsDir, entry);
        const indexPath = join(entryPath, "index.ts");

        if (statSync(entryPath).isDirectory() && existsSync(indexPath)) {
            components.push(entry);
            toolToComponentMap[toKebabCase(entry)] = entry;
        }
    }

    return { components, toolToComponentMap };
}

const { components, toolToComponentMap } = discoverComponents();

/**
 * Vite plugin that generates HTML entry files for each discovered component
 * based on the template.html file.
 */
function generateHtmlEntries(): Plugin {
    return {
        name: "generate-html-entries",
        buildStart() {
            const template = readFileSync(templatePath, "utf-8");

            if (!existsSync(entriesDir)) {
                mkdirSync(entriesDir, { recursive: true });
            }

            for (const componentName of components) {
                const html = template
                    .replace("{{COMPONENT_NAME}}", componentName)
                    .replace("{{TITLE}}", componentName.replace(/([A-Z])/g, " $1").trim()) // "ListDatabases" -> "List Databases"
                    .replace("{{MOUNT_PATH}}", mountPath);

                const outputPath = join(entriesDir, `${componentName}.html`);
                writeFileSync(outputPath, html);
                console.log(`[generate-html-entries] Generated ${componentName}.html`);
            }
        },
    };
}

/**
 * Vite plugin that generates per-tool UI modules after the build completes.
 */
function generateUIModule(): Plugin {
    return {
        name: "generate-ui-module",
        closeBundle() {
            if (!existsSync(uiDistPath)) {
                console.warn("[generate-ui-module] dist/ui not found, skipping module generation");
                return;
            }

            const toolsDir = join(generatedDir, "tools");
            mkdirSync(toolsDir, { recursive: true });
            const existingToolFiles = readdirSync(toolsDir).filter((file) => file.endsWith(".ts"));

            const generatedTools: string[] = [];

            for (const [toolName, componentName] of Object.entries(toolToComponentMap)) {
                const htmlFile = join(uiDistPath, `${componentName}.html`);
                if (!existsSync(htmlFile)) {
                    console.warn(
                        `[generate-ui-module] HTML file not found for component "${componentName}" (tool: "${toolName}")`
                    );
                    continue;
                }
                const html = readFileSync(htmlFile, "utf-8");

                const toolModuleContent = `/**
 * AUTO-GENERATED FILE - DO NOT EDIT MANUALLY
 * Generated by: vite build --config vite.ui.config.ts
 * Tool: ${toolName}
 * Component: ${componentName}
 */
export const ${componentName}Html = ${JSON.stringify(html)};
`;
                writeFileSync(join(toolsDir, `${toolName}.ts`), toolModuleContent);
                generatedTools.push(toolName);
            }

            // Generate the loaders.ts file with lazy import functions for each tool
            // Uses .js extension for ESM compatibility (tsc compiles .ts -> .js)
            const loaderEntries = generatedTools
                .map((toolName) => {
                    const componentName = toolToComponentMap[toolName];
                    return `    "${toolName}": async () => {
        const mod = await import("./tools/${toolName}.js");
        return mod.${componentName}Html;
    }`;
                })
                .join(",\n");

            const loadersContent = `/**
 * AUTO-GENERATED FILE - DO NOT EDIT MANUALLY
 * Generated by: pnpm generate:ui
 *
 * Lazy loaders for UI modules. Each loader returns a Promise<string> with the HTML.
 */
export const uiLoaders: Record<string, () => Promise<string>> = {
${loaderEntries}
};
`;
            writeFileSync(join(generatedDir, "loaders.ts"), loadersContent);

            console.log(
                `[generate-ui-module] Generated ${generatedTools.length} lazy UI module(s): ${generatedTools.join(", ")}`
            );
            console.log(`[generate-ui-module] Generated loaders.ts with ${generatedTools.length} loader(s)`);

            // Remove stale tool modules from previous builds (e.g., when a UI component was deleted)
            const staleTools = existingToolFiles.filter((file) => {
                const toolName = file.replace(/\.ts$/, "");
                return !generatedTools.includes(toolName);
            });
            for (const staleTool of staleTools) {
                rmSync(join(toolsDir, staleTool));
                console.log(`[generate-ui-module] Removed stale tool module: ${staleTool}`);
            }

            // Clean up intermediate dist/ui directory - the HTML is now embedded in the .ts modules
            if (existsSync(uiDistPath)) {
                rmSync(uiDistPath, { recursive: true });
                console.log("[generate-ui-module] Cleaned up intermediate dist/ui directory");
            }
        },
    };
}

export default defineConfig({
    root: entriesDir,
    plugins: [
        generateHtmlEntries(),
        nodePolyfills({
            include: ["buffer", "stream"],
            globals: {
                Buffer: true,
            },
        }) as unknown as PluginOption,
        react(),
        viteSingleFile({
            removeViteModuleLoader: true,
        }),
        generateUIModule(),
    ],
    build: {
        outDir: resolve(__dirname, "dist/ui"),
        emptyOutDir: true,
        rollupOptions: {
            input: Object.fromEntries(components.map((name) => [name, resolve(entriesDir, `${name}.html`)])),
        },
        assetsInlineLimit: 100000000,
        sourcemap: false,
        minify: "esbuild",
    },
    resolve: {
        alias: {
            "@ui": resolve(__dirname, "src/ui"),
        },
    },
});