import * as t from "@babel/types";
import { getStableEntityId, SOURCE_ID_ATTRIBUTE, } from "@superblocksteam/library-shared";
import { generatePredictableId } from "../ids.js";
import { traverse } from "../traverse.js";
import { makeJSXAttribute } from "./jsx.js";
import { getTagName, isSbScopeFunction, nodeToValue } from "./util.js";
// TODO(code-mode): profile if a two-pass approach is too slow for actual usage
// this makes generating stable ids nice and easy, though
export function supplementElementIds({ fileName, ast, shouldModifyAst, scopeId, }) {
    const tagCount = new Map();
    const idMap = new Map();
    traverse(ast, {
        JSXOpeningElement(path) {
            if (path.node.name.type === "JSXIdentifier" ||
                path.node.name.type === "JSXMemberExpression" ||
                path.node.name.type === "JSXNamespacedName") {
                const tagName = getTagName(path.node.name);
                if (tagName) {
                    const parentPath = generateParentPath(path);
                    const count = tagCount.get(tagName) ?? 0;
                    const id = getSbElementId(path.node) ??
                        `SB-${generatePredictableId(fileName + parentPath.join("-") + count)}`;
                    idMap.set(id, path);
                    if (shouldModifyAst) {
                        path.pushContainer("attributes", makeJSXAttribute(SOURCE_ID_ATTRIBUTE, id));
                    }
                    // This is used to persist IDs so the Vite plugin can read them
                    setSbElementId(path.node, id);
                    tagCount.set(tagName, count + 1);
                }
            }
        },
        CallExpression(path) {
            // Check for createScope calls
            const isCreateScope = path.type === "CallExpression" &&
                path.node.callee.type === "Identifier" &&
                isSbScopeFunction(path.node.callee.name);
            if (!isCreateScope) {
                return;
            }
            // For createScope, the first argument is now a function that returns an object
            if (scopeId && path.node.arguments.length > 0) {
                const firstArg = path.get("arguments.0");
                // Handle function-based scope: createScope(() => ({ ... }))
                if (!Array.isArray(firstArg) &&
                    (firstArg.isArrowFunctionExpression() ||
                        firstArg.isFunctionExpression())) {
                    const functionBody = firstArg.get("body");
                    let entitiesObject = null;
                    if (!Array.isArray(functionBody) &&
                        functionBody.isObjectExpression()) {
                        // Arrow function with expression body: () => ({ ... })
                        entitiesObject = functionBody;
                    }
                    else if (!Array.isArray(functionBody) &&
                        functionBody.isBlockStatement()) {
                        // Function with block body: () => { return { ... }; }
                        functionBody.traverse({
                            ReturnStatement(returnPath) {
                                const argument = returnPath.get("argument");
                                if (argument.isObjectExpression()) {
                                    entitiesObject = argument;
                                }
                            },
                        });
                    }
                    // Process the entities object if found
                    if (entitiesObject && t.isObjectExpression(entitiesObject.node)) {
                        entitiesObject.node.properties.forEach((property, index) => {
                            if (property.type === "ObjectProperty") {
                                const name = nodeToValue(property.key);
                                const id = getStableEntityId(scopeId, name);
                                idMap.set(id, entitiesObject.get(`properties.${index}`));
                                setSbElementId(property, id);
                                if (shouldModifyAst) {
                                    if (t.isCallExpression(property.value) &&
                                        t.isObjectExpression(property.value.arguments[0])) {
                                        const entityConfig = property.value.arguments[0];
                                        entityConfig.properties.push(t.objectProperty(t.identifier("id"), t.stringLiteral(id)));
                                    }
                                }
                            }
                        });
                    }
                }
            }
            // Also keep track of component types defined in the createScope call
            // we use a different id for these (than for the actual component)
            if (scopeId && path.node.typeParameters) {
                const typeParams = path.node.typeParameters.params;
                if (typeParams.length > 0 && typeParams[0]?.type === "TSTypeLiteral") {
                    const typeMembers = typeParams[0].members;
                    typeMembers.forEach((member, index) => {
                        if (member.type === "TSPropertySignature" && member.key) {
                            let componentName = null;
                            // Extract component name from the key
                            if (member.key.type === "Identifier") {
                                componentName = member.key.name;
                            }
                            else if (member.key.type === "StringLiteral") {
                                componentName = member.key.value;
                            }
                            if (componentName) {
                                const id = `SB-${generatePredictableId(scopeId + componentName)}-type`;
                                // Get the path to this type member
                                const typeParamsPath = path.get("typeParameters.params.0.members");
                                if (typeParamsPath && typeParamsPath[index]) {
                                    idMap.set(id, typeParamsPath[index]);
                                }
                            }
                        }
                    });
                }
            }
        },
    });
    return { idMap };
}
export function setSbElementId(node, id) {
    node.extra ??= {};
    node.extra.sbElementId = id;
}
export function getSbElementId(node) {
    return node.extra?.sbElementId;
}
function generateParentPath(path) {
    const parentPath = [];
    let currentPath = path;
    while (currentPath.parentPath) {
        if (currentPath.type === "JSXElement") {
            const tagName = getTagName(currentPath.node);
            if (tagName) {
                parentPath.push(tagName);
            }
        }
        currentPath = currentPath.parentPath;
    }
    return parentPath;
}
//# sourceMappingURL=ids.js.map