penpot/mcp/packages/plugin/src/PenpotUtils.ts

import type {
    Board,
    Bounds,
    Fill,
    FlexLayout,
    GridLayout,
    Page,
    Rectangle,
    Shape,
    Text,
    VariantContainer,
} from "@penpot/plugin-types";

export class PenpotUtils {
    /**
     * Generates an overview structure of the given shape,
     * providing its id, name and type, and recursively its children's attributes.
     * The `type` field indicates the type in the Penpot API.
     * If the shape has a layout system (flex or grid), includes layout information.
     *
     * @param shape - The root shape to generate the structure from
     * @param maxDepth - Optional maximum depth to traverse (leave undefined for unlimited)
     * @returns An object representing the shape structure
     */
    public static shapeStructure(shape: Shape, maxDepth: number | undefined = undefined): object {
        let children = undefined;
        if (maxDepth === undefined || maxDepth > 0) {
            if ("children" in shape && shape.children) {
                children = shape.children.map((child) =>
                    this.shapeStructure(child, maxDepth === undefined ? undefined : maxDepth - 1)
                );
            }
        }

        const result: any = {
            id: shape.id,
            name: shape.name,
            type: shape.type,
        };

        // add layout information if present
        if ("flex" in shape && shape.flex) {
            const flex: FlexLayout = shape.flex;
            result.layout = {
                type: "flex",
                dir: flex.dir,
                rowGap: flex.rowGap,
                columnGap: flex.columnGap,
            };
        } else if ("grid" in shape && shape.grid) {
            const grid: GridLayout = shape.grid;
            result.layout = {
                type: "grid",
                rows: grid.rows,
                columns: grid.columns,
                rowGap: grid.rowGap,
                columnGap: grid.columnGap,
            };
        }

        // add component instance information if present
        if (shape.isComponentInstance()) {
            result.componentInstance = {};
            const component = shape.component();
            if (component) {
                result.componentInstance.componentId = component.id;
                result.componentInstance.componentName = component.name;
                const mainInstance = component.mainInstance();
                if (mainInstance) {
                    result.componentInstance.mainInstanceId = mainInstance.id;
                }
            }
        }

        // finally, add children (last for more readable nesting order)
        result.children = children;

        return result;
    }

    /**
     * Finds all shapes that matches the given predicate in the given shape tree.
     *
     * @param predicate - A function that takes a shape and returns true if it matches the criteria
     * @param root - The root shape to start the search from (if null, searches all pages)
     */
    public static findShapes(predicate: (shape: Shape) => boolean, root: Shape | null = null): Shape[] {
        let result = new Array<Shape>();

        // Walk the descendants of the given shape, testing the predicate on each
        // one. The starting shape is never tested itself: findShapes searches
        // *within* a container, so a predicate meant for descendants cannot
        // accidentally select (and later mutate) the container it was handed.
        let walk = function (shape: Shape | null) {
            if (!shape) {
                return;
            }
            if ("children" in shape && shape.children) {
                for (let child of shape.children) {
                    if (predicate(child)) {
                        result.push(child);
                    }
                    walk(child);
                }
            }
        };

        if (root === null) {
            const pages = penpot.currentFile?.pages;
            if (pages) {
                for (let page of pages) {
                    walk(page.root);
                }
            }
        } else {
            walk(root);
        }
        return result;
    }

    /**
     * Finds the first shape that matches the given predicate in the given shape tree.
     *
     * @param predicate - A function that takes a shape and returns true if it matches the criteria
     * @param root - The root shape to start the search from (if null, searches all pages)
     */
    public static findShape(predicate: (shape: Shape) => boolean, root: Shape | null = null): Shape | null {
        let find = function (shape: Shape | null): Shape | null {
            if (!shape) {
                return null;
            }
            if (predicate(shape)) {
                return shape;
            }
            if ("children" in shape && shape.children) {
                for (let child of shape.children) {
                    let result = find(child);
                    if (result) {
                        return result;
                    }
                }
            }
            return null;
        };

        if (root === null) {
            const pages = penpot.currentFile?.pages;
            if (pages) {
                for (let page of pages) {
                    let result = find(page.root);
                    if (result) {
                        return result;
                    }
                }
            }
            return null;
        } else {
            return find(root);
        }
    }

    /**
     * Finds a shape by its unique ID.
     *
     * @param id - The unique ID of the shape to find
     * @returns The shape with the matching ID, or null if not found
     */
    public static findShapeById(id: string): Shape | null {
        return this.findShape((shape) => shape.id === id);
    }

    public static findPage(predicate: (page: Page) => boolean): Page | null {
        let page = penpot.currentFile!.pages.find(predicate);
        return page || null;
    }

    public static getPages(): { id: string; name: string }[] {
        return penpot.currentFile!.pages.map((page) => ({ id: page.id, name: page.name }));
    }

    public static getPageById(id: string): Page | null {
        return this.findPage((page) => page.id === id);
    }

    public static getPageByName(name: string): Page | null {
        return this.findPage((page) => page.name.toLowerCase() === name.toLowerCase());
    }

    public static getPageForShape(shape: Shape): Page | null {
        for (const page of penpot.currentFile!.pages) {
            if (page.getShapeById(shape.id)) {
                return page;
            }
        }
        return null;
    }

    public static generateCss(shape: Shape): string {
        const page = this.getPageForShape(shape);
        if (!page) {
            throw new Error("Shape is not part of any page");
        }
        penpot.openPage(page);
        return penpot.generateStyle([shape], { type: "css", includeChildren: true });
    }

    /**
     * Gets the actual rendering bounds of a shape. For most shapes, this is simply the `bounds` property.
     * However, for Text shapes, the `bounds` may not reflect the true size of the rendered text content,
     * so we use the `textBounds` property instead.
     *
     * @param shape - The shape to get the bounds for
     */
    public static getBounds(shape: Shape): Bounds {
        if (shape.type === "text") {
            const text = shape as Text;
            return text.textBounds;
        } else {
            return shape.bounds;
        }
    }

    /**
     * Checks if a child shape is fully contained within its parent's bounds.
     * Visual containment means all edges of the child are within the parent's bounding box.
     *
     * @param child - The child shape to check
     * @param parent - The parent shape to check against
     * @returns true if child is fully contained within parent bounds, false otherwise
     */
    public static isContainedIn(child: Shape, parent: Shape): boolean {
        const childBounds = this.getBounds(child);
        const parentBounds = this.getBounds(parent);
        return (
            childBounds.x >= parentBounds.x &&
            childBounds.y >= parentBounds.y &&
            childBounds.x + childBounds.width <= parentBounds.x + parentBounds.width &&
            childBounds.y + childBounds.height <= parentBounds.y + parentBounds.height
        );
    }

    /**
     * Sets the position of a shape relative to its parent's position.
     * This is a convenience method since parentX and parentY are read-only properties.
     *
     * @param shape - The shape to position
     * @param parentX - The desired X position relative to the parent
     * @param parentY - The desired Y position relative to the parent
     * @throws Error if the shape has no parent
     */
    public static setParentXY(shape: Shape, parentX: number, parentY: number): void {
        if (!shape.parent) {
            throw new Error("Shape has no parent - cannot set parent-relative position");
        }
        shape.x = shape.parent.x + parentX;
        shape.y = shape.parent.y + parentY;
    }

    /**
     * Adds a flex layout to a container while preserving the visual order of existing children.
     * Without this, adding a flex layout can arbitrarily reorder children.
     *
     * The method sorts children by their current position (x for "row", y for "column") before
     * adding the layout, then reorders them to maintain that visual sequence.
     *
     * @param container - The container (board) to add the flex layout to
     * @param dir - The layout direction: "row" for horizontal, "column" for vertical
     * @returns The created FlexLayout instance
     */
    public static addFlexLayout(container: Board, dir: "column" | "row"): FlexLayout {
        // obtain children sorted by position (ascending)
        const children = "children" in container && container.children ? [...container.children] : [];
        const sortedChildren = children.sort((a, b) => (dir === "row" ? a.x - b.x : a.y - b.y));

        // add the flex layout
        const flexLayout = container.addFlexLayout();
        flexLayout.dir = dir;

        // reorder children to preserve visual order; since the children array is reversed
        // relative to visual order for dir="column" or dir="row", we insert each child at
        // index 0 in sorted order, which places the first (smallest position) at the highest
        // index, making it appear first visually
        for (const child of sortedChildren) {
            child.setParentIndex(0);
        }

        return flexLayout;
    }

    /**
     * Analyzes all descendants of a shape by applying an evaluator function to each.
     * Only descendants for which the evaluator returns a non-null/non-undefined value are included in the result.
     * This is a general-purpose utility for validation, analysis, or collecting corrector functions.
     *
     * @param root - The root shape whose descendants to analyze
     * @param evaluator - Function called for each descendant with (root, descendant); return null/undefined to skip
     * @param maxDepth - Optional maximum depth to traverse (undefined for unlimited)
     * @returns Array of objects containing the shape and the evaluator's result
     */
    public static analyzeDescendants<T>(
        root: Shape,
        evaluator: (root: Shape, descendant: Shape) => T | null | undefined,
        maxDepth: number | undefined = undefined
    ): Array<{ shape: Shape; result: NonNullable<T> }> {
        const results: Array<{ shape: Shape; result: NonNullable<T> }> = [];

        const traverse = (shape: Shape, currentDepth: number): void => {
            const result = evaluator(root, shape);
            if (result !== null && result !== undefined) {
                results.push({ shape, result: result as NonNullable<T> });
            }

            if (maxDepth === undefined || currentDepth < maxDepth) {
                if ("children" in shape && shape.children) {
                    for (const child of shape.children) {
                        traverse(child, currentDepth + 1);
                    }
                }
            }
        };

        // Start traversal with root's children (not root itself)
        if ("children" in root && root.children) {
            for (const child of root.children) {
                traverse(child, 1);
            }
        }

        return results;
    }

    /**
     * Decodes a base64 string to a Uint8Array.
     *
     * @param base64 - The base64-encoded string to decode
     * @returns The decoded data as a Uint8Array
     */
    public static base64ToByteArray(base64: string): Uint8Array {
        const binary = atob(base64);
        const bytes = new Uint8Array(binary.length);
        for (let i = 0; i < binary.length; i++) {
            bytes[i] = binary.charCodeAt(i);
        }
        return bytes;
    }

    /**
     * Imports an image from base64 data into the Penpot design as a Rectangle shape filled with the image.
     * The rectangle has the image's original proportions by default.
     * Optionally accepts position (x, y) and dimensions (width, height) parameters.
     * If only one dimension is provided, the other is calculated to maintain the image's aspect ratio.
     *
     * This function is used internally by the ImportImageTool in the MCP server.
     *
     * @param base64 - The base64-encoded image data
     * @param mimeType - The MIME type of the image (e.g., "image/png")
     * @param name - The name to assign to the newly created rectangle shape
     * @param x - The x-coordinate for positioning the rectangle (optional)
     * @param y - The y-coordinate for positioning the rectangle (optional)
     * @param width - The desired width of the rectangle (optional)
     * @param height - The desired height of the rectangle (optional)
     */
    public static async importImage(
        base64: string,
        mimeType: string,
        name: string,
        x: number | undefined,
        y: number | undefined,
        width: number | undefined,
        height: number | undefined
    ): Promise<Rectangle> {
        // convert base64 to Uint8Array
        const bytes = PenpotUtils.base64ToByteArray(base64);

        // upload the image data to Penpot
        const imageData = await penpot.uploadMediaData(name, bytes, mimeType);

        // create a rectangle shape
        const rect = penpot.createRectangle();
        rect.name = name;

        // calculate dimensions
        let rectWidth, rectHeight;
        const hasWidth = width !== undefined;
        const hasHeight = height !== undefined;

        if (hasWidth && hasHeight) {
            // both width and height provided - use them directly
            rectWidth = width;
            rectHeight = height;
        } else if (hasWidth) {
            // only width provided - maintain aspect ratio
            rectWidth = width;
            rectHeight = rectWidth * (imageData.height / imageData.width);
        } else if (hasHeight) {
            // only height provided - maintain aspect ratio
            rectHeight = height;
            rectWidth = rectHeight * (imageData.width / imageData.height);
        } else {
            // neither provided - use original dimensions
            rectWidth = imageData.width;
            rectHeight = imageData.height;
        }

        // set rectangle dimensions
        rect.resize(rectWidth, rectHeight);

        // set position if provided
        if (x !== undefined) {
            rect.x = x;
        }
        if (y !== undefined) {
            rect.y = y;
        }

        // apply the image as a fill
        rect.fills = [{ fillOpacity: 1, fillImage: imageData }];

        return rect;
    }

    /**
     * Exports the given shape (or its fill) to BASE64 image data.
     *
     * This function is used internally by the ExportImageTool in the MCP server.
     *
     * @param shape - The shape whose image data to export
     * @param mode - Either "shape" (to export the entire shape, including descendants) or "fill"
     *    to export the shape's raw fill image data
     * @param asSVG - Whether to export as SVG rather than as a pixel image (only supported for mode "shape")
     * @returns A byte array containing the exported image data.
     *   - For mode="shape", it will be PNG or SVG data depending on the value of `asSVG`.
     *   - For mode="fill", it will be whatever format the fill image is stored in.
     */
    public static async exportImage(shape: Shape, mode: "shape" | "fill", asSVG: boolean): Promise<Uint8Array> {
        // Updates are asynchronous in Penpot, so wait a tick to ensure any pending updates are applied before export.
        // The constant wait time is a temporary workardound until a better solution for penpot/penpot-mcp#27
        // is implemented.
        await new Promise((resolve) => setTimeout(resolve, 200));
        // Perform export
        switch (mode) {
            case "shape":
                return shape.export({ type: asSVG ? "svg" : "png" });
            case "fill":
                if (asSVG) {
                    throw new Error("Image fills cannot be exported as SVG");
                }
                // check whether the shape has the `fills` member
                if (!("fills" in shape)) {
                    throw new Error("Shape with `fills` member is required for fill export mode");
                }
                // find first fill that has fillImage
                const fills: Fill[] = (shape as any).fills;
                for (const fill of fills) {
                    if (fill.fillImage) {
                        const imageData = fill.fillImage;
                        return imageData.data();
                    }
                }
                throw new Error("No fill with image data found in the shape");
            default:
                throw new Error(`Unsupported export mode: ${mode}`);
        }
    }

    /**
     * Finds all tokens that match the given name across all token sets.
     *
     * @param name - The name of the token to search for (case-sensitive exact match)
     * @returns An array of all matching tokens (may be empty)
     */
    public static findTokensByName(name: string): any[] {
        const tokens: any[] = [];
        const tokenCatalog = penpot.library.local.tokens;

        for (const set of tokenCatalog.sets) {
            for (const token of set.tokens) {
                if (token.name === name) {
                    tokens.push(token);
                }
            }
        }

        return tokens;
    }

    /**
     * Finds the first token that matches the given name across all token sets.
     *
     * @param name - The name of the token to search for (case-sensitive exact match)
     * @returns The first matching token, or null if not found
     */
    public static findTokenByName(name: string): any | null {
        const tokenCatalog = penpot.library.local.tokens;

        for (const set of tokenCatalog.sets) {
            for (const token of set.tokens) {
                if (token.name === name) {
                    return token;
                }
            }
        }

        return null;
    }

    /**
     * Gets the token set that contains the given token.
     *
     * @param token - The token whose set to find
     * @returns The TokenSet containing this token, or null if not found
     */
    public static getTokenSet(token: any): any | null {
        const tokenCatalog = penpot.library.local.tokens;

        for (const set of tokenCatalog.sets) {
            if (set.tokens.includes(token)) {
                return set;
            }
        }

        return null;
    }

    /**
     * Creates a variant group from a list of main component instances and immediately
     * configures its property names and values.
     *
     * Encapsulates the full multi-step workflow in a single call:
     *   1. `penpot.createVariantFromComponents()` — groups the components
     *   2. `variants.renameProperty()` / `variants.addProperty()` — sets property names
     *   3. `comp.setVariantProperty()` — sets each component's property values
     *
     * @param components - Array of main-instance boards to combine. Each element carries
     *   a `shape` (the Board on the canvas) and a `properties` map of
     *   `{ propertyName: value }` for that component.
     * @returns The newly created `VariantContainer`.
     *
     * @example
     * // Three button variants differing in Size
     * // Given three main components s, m, l (here: the first three main components on the page)
     * const [s, m, l] = penpot.currentPage.findAllShapes(sh => sh.isMainComponent()).slice(0, 3) as Board[];
     * const container = PenpotUtils.createVariantContainer([
     *   { shape: s, properties: { Size: 'Small' } },
     *   { shape: m, properties: { Size: 'Medium' } },
     *   { shape: l, properties: { Size: 'Large' } },
     * ]);
     *
     * @example
     * // Two properties: Size × State
     * const container = PenpotUtils.createVariantContainer([
     *   { shape: s, properties: { Size: 'Small', State: 'Default' } },
     *   { shape: m, properties: { Size: 'Medium', State: 'Default' } },
     *   { shape: l, properties: { Size: 'Large', State: 'Hover' } },
     * ]);
     */
    public static createVariantContainer(
        components: Array<{ shape: Board; properties: Record<string, string> }>
    ): VariantContainer {
        // Collect all unique property names (preserving first-seen order)
        const propNames: string[] = [];
        for (const { properties } of components) {
            for (const name of Object.keys(properties)) {
                if (!propNames.includes(name)) propNames.push(name);
            }
        }

        // 1. Create the variant container
        const container: VariantContainer = (penpot as any).createVariantFromComponents(components.map((c) => c.shape));
        const variants = container.variants;
        if (!variants) {
            throw new Error(
                "createVariantContainer: the created container has no `variants`. " +
                    "Ensure every provided shape is a main component instance (Board)."
            );
        }

        // 2. Rename / add properties
        // createVariantFromComponents always creates exactly one property ("Property 1")
        for (let i = 0; i < propNames.length; i++) {
            if (i === 0) {
                variants.renameProperty(0, propNames[0]);
            } else {
                variants.addProperty();
                variants.renameProperty(i, propNames[i]);
            }
        }

        // 3. Set each component's property values
        const variantComps = variants.variantComponents();
        for (let compIdx = 0; compIdx < variantComps.length; compIdx++) {
            const comp = variantComps[compIdx];
            if (!penpot.utils.types.isVariantComponent(comp)) continue;
            const props = components[compIdx]?.properties ?? {};
            for (let propIdx = 0; propIdx < propNames.length; propIdx++) {
                const value = props[propNames[propIdx]];
                if (value !== undefined) {
                    comp.setVariantProperty(propIdx, value);
                }
            }
        }

        return container;
    }

    /**
     * Generates an overview of all tokens organized by token set name, token type, and token name.
     * The result is a nested object structure: {tokenSetName: {tokenType: [tokenName, ...]}}.
     *
     * @returns An object mapping token set names to objects that map token types to arrays of token names
     */
    public static tokenOverview(): Record<string, Record<string, string[]>> {
        const overview: Record<string, Record<string, string[]>> = {};
        const tokenCatalog = penpot.library.local.tokens;

        for (const set of tokenCatalog.sets) {
            const setOverview: Record<string, string[]> = {};

            for (const token of set.tokens) {
                const tokenType = token.type;
                if (!setOverview[tokenType]) {
                    setOverview[tokenType] = [];
                }
                setOverview[tokenType].push(token.name);
            }

            overview[set.name] = setOverview;
        }

        return overview;
    }
}