module.exports = function addDataLocator({ types: t }) {
  function getElementName(node) {
    if (!node) return "";
    if (t.isJSXIdentifier(node)) {
      return node.name;
    }
    if (t.isJSXMemberExpression(node)) {
      return `${getElementName(node.object)}.${node.property.name}`;
    }
    return "";
  }

  const FRAGMENT_DEFAULT_COMPONENTS = new Set([
    "Fragment",
    "React.Fragment",
    "Tab.Group",
    "Tabs",
    "Transition",
    "Transition.Root",
    "Transition.Child",
    "Menu",
    "Menu.Item",
    "Switch.Group",
    "Popover.Group",
  ]);

  return {
    name: "add-data-locator",
    visitor: {
      JSXOpeningElement(path, state) {
        if (process.env.NODE_ENV && process.env.NODE_ENV !== "development") {
          return;
        }

        const filePath = state.file.opts.filename;

        if (!filePath || filePath.includes("node_modules")) {
          return;
        }

        const elementName = getElementName(path.node.name);

        // 1. Skip React Fragment and known Headless UI components that render Fragment
        if (
          !elementName ||
          elementName === "Fragment" ||
          elementName.endsWith(".Fragment") ||
          FRAGMENT_DEFAULT_COMPONENTS.has(elementName)
        ) {
          return;
        }

        // 2. Skip any component explicitly rendered as a Fragment (e.g. as={Fragment} or as="Fragment")
        const hasFragmentAs = path.node.attributes.some((attribute) => {
          if (
            t.isJSXAttribute(attribute) &&
            t.isJSXIdentifier(attribute.name) &&
            attribute.name.name === "as"
          ) {
            if (t.isJSXExpressionContainer(attribute.value)) {
              const expr = attribute.value.expression;
              if (t.isIdentifier(expr, { name: "Fragment" })) return true;
              if (
                t.isMemberExpression(expr) &&
                t.isIdentifier(expr.property, { name: "Fragment" })
              ) {
                return true;
              }
            }
            if (
              t.isStringLiteral(attribute.value) &&
              attribute.value.value === "Fragment"
            ) {
              return true;
            }
          }
          return false;
        });

        if (hasFragmentAs) {
          return;
        }

        const attributeExists = path.node.attributes.some(
          (attribute) =>
            t.isJSXAttribute(attribute) &&
            t.isJSXIdentifier(attribute.name) &&
            attribute.name.name === "data-locator",
        );

        if (attributeExists) {
          return;
        }

        const lineNumber = path.node.loc?.start.line ?? "unknown";
        const columnNumber = path.node.loc?.start.column ?? "unknown";
        const locatorValue = `${filePath}:${lineNumber}:${columnNumber}`;

        path.node.attributes.push(
          t.jsxAttribute(
            t.jsxIdentifier("data-locator"),
            t.stringLiteral(locatorValue),
          ),
        );
      },
    },
  };
};