nextjs-ide-helper 1.1.2 → 1.2.0

package/README.md

@@ -12,6 +12,8 @@ A Next.js plugin that automatically adds IDE buttons to React components in deve
 - ⚡ **Hydration Safe**: No SSR/client hydration mismatches
 - 📱 **TypeScript Support**: Full TypeScript definitions included
 - 🔌 **Multi-IDE Support**: Supports Cursor, VS Code, WebStorm, and Atom
+- 🎭 **Multiple Component Patterns**: Supports all React component export patterns
+- 🔍 **AST-Based Processing**: Uses robust Abstract Syntax Tree parsing for accurate code transformation
 
 ## Installation
 
@@ -91,16 +93,95 @@ export default withIdeButton(MyComponent, 'src/components/MyComponent.tsx', {
 });
 ```
 
+## Supported Component Patterns
+
+The plugin automatically detects and wraps all types of React component export patterns:
+
+### Named Components
+```tsx
+// Standard pattern
+const MyComponent = () => <div>Hello</div>;
+export default MyComponent;
+
+// Variable declaration
+const MyComponent = function() {
+  return <div>Hello</div>;
+};
+export default MyComponent;
+```
+
+### Direct Export Function Components
+```tsx
+// Named function
+export default function MyComponent() {
+  return <div>Hello</div>;
+}
+
+// Anonymous function
+export default function() {
+  return <div>Hello</div>;
+}
+```
+
+### Arrow Function Components
+```tsx
+// Anonymous arrow function
+export default () => {
+  return <div>Hello</div>;
+};
+
+// Arrow function with parameters
+export default (props) => {
+  return <div>Hello {props.name}</div>;
+};
+```
+
+### Class Components
+```tsx
+// Named class
+export default class MyComponent extends React.Component {
+  render() {
+    return <div>Hello</div>;
+  }
+}
+
+// TypeScript class
+export default class MyComponent extends Component<Props> {
+  render() {
+    return <div>Hello</div>;
+  }
+}
+```
+
+### TypeScript Components
+```tsx
+// Function with TypeScript
+interface Props {
+  title: string;
+}
+
+export default function MyComponent(props: Props) {
+  return <div>{props.title}</div>;
+}
+
+// Arrow function with TypeScript
+export default (props: Props) => {
+  return <div>{props.title}</div>;
+};
+```
+
 ## How It Works
 
-1. **Webpack Loader**: The plugin uses a custom webpack loader to transform your React components at build time
-2. **Automatic Detection**: It scans specified directories for `.tsx` and `.jsx` files
-3. **Smart Wrapping**: Only wraps components that export a default React component
-4. **Development Only**: The IDE buttons only appear in development mode
-5. **Hydration Safe**: Uses client-side state to prevent SSR/hydration mismatches
+1. **AST-Based Processing**: Uses Babel's Abstract Syntax Tree parser for robust code analysis
+2. **Webpack Loader**: The plugin uses a custom webpack loader to transform your React components at build time
+3. **Automatic Detection**: It scans specified directories for `.tsx` and `.jsx` files
+4. **Smart Wrapping**: Only wraps components that export a default React component (PascalCase names)
+5. **Development Only**: The IDE buttons only appear in development mode
+6. **Hydration Safe**: Uses client-side state to prevent SSR/hydration mismatches
 
-## Example
+## Examples
 
+### Standard Component Pattern
 Before (your original component):
 ```tsx
 // src/components/Button.tsx
@@ -133,7 +214,53 @@ const Button = ({ children, onClick }) => {
 
 export default withIdeButton(Button, 'src/components/Button.tsx', { 
   projectRoot: '/path/to/project',
-  ideType: 'cursor' // defaults to cursor, can be 'vscode', 'webstorm', 'atom'
+  ideType: 'cursor'
+});
+```
+
+### Direct Export Function Component
+Before:
+```tsx
+// src/components/Header.tsx
+export default function Header() {
+  return <header>My App Header</header>;
+}
+```
+
+After (automatically transformed):
+```tsx
+// src/components/Header.tsx (transformed by the plugin)
+import { withIdeButton } from 'nextjs-ide-plugin/withIdeButton';
+
+function Header() {
+  return <header>My App Header</header>;
+}
+
+export default withIdeButton(Header, 'src/components/Header.tsx', {
+  projectRoot: '/path/to/project',
+  ideType: 'cursor'
+});
+```
+
+### Anonymous Component
+Before:
+```tsx
+// src/components/Footer.tsx
+export default () => {
+  return <footer>© 2025 My App</footer>;
+};
+```
+
+After (automatically transformed):
+```tsx
+// src/components/Footer.tsx (transformed by the plugin)
+import { withIdeButton } from 'nextjs-ide-plugin/withIdeButton';
+
+export default withIdeButton(() => {
+  return <footer>© 2025 My App</footer>;
+}, 'src/components/Footer.tsx', {
+  projectRoot: '/path/to/project',
+  ideType: 'cursor'
 });
 ```
 
@@ -174,8 +301,24 @@ MIT
 
 ## Changelog
 
+See [CHANGELOG.md](./CHANGELOG.md) for detailed release notes.
+
+### Recent Releases
+
+### 1.2.0 - Enhanced Component Support
+- Added support for all React component export patterns
+- AST-based code transformation for better reliability
+- Anonymous component support
+- Enhanced TypeScript compatibility
+
+### 1.1.3
+- Updated readme
+
+### 1.1.2
+- Added support for multiple IDEs (Cursor, VS Code, WebStorm, Atom)
+
 ### 1.0.0
 - Initial release
 - Automatic component wrapping
 - TypeScript support
-- Hydration-safe implementation
+- Hydration-safe implementation

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "nextjs-ide-helper",
-  "version": "1.1.2",
+  "version": "1.2.0",
   "description": "A Next.js plugin that automatically adds IDE buttons to React components for seamless IDE integration. Supports Cursor, VS Code, WebStorm, and Atom.",
   "main": "lib/index.js",
   "types": "lib/index.d.ts",
@@ -11,7 +11,7 @@
   "scripts": {
     "build": "tsc",
     "prepare": "npm run build",
-    "test": "echo \"Error: no test specified\" && exit 1"
+    "test": "jest"
   },
   "keywords": [
     "nextjs",
@@ -33,8 +33,14 @@
     "react": ">=18.0.0"
   },
   "devDependencies": {
+    "@babel/generator": "^7.28.0",
+    "@babel/parser": "^7.28.0",
+    "@babel/traverse": "^7.28.0",
+    "@babel/types": "^7.28.2",
+    "@types/jest": "^30.0.0",
     "@types/node": "^20.19.9",
     "@types/react": "^18.3.23",
+    "jest": "^30.0.5",
     "typescript": "^5.8.3"
   },
   "repository": {

package/src/__tests__/loader.test.js

@@ -0,0 +1,429 @@
+const path = require('path');
+const loader = require('../loader.js');
+
+describe('cursorButtonLoader', function() {
+  let mockContext;
+
+  beforeEach(function() {
+    mockContext = {
+      resourcePath: process.cwd() + '/src/components/Button.tsx',
+      getOptions: jest.fn(() => ({}))
+    };
+    
+    process.env.NODE_ENV = 'development';
+  });
+
+  afterEach(function() {
+    jest.clearAllMocks();
+  });
+
+  describe('when enabled', function() {
+    beforeEach(function() {
+      mockContext.getOptions.mockReturnValue({ enabled: true });
+    });
+
+    it('should wrap a React component with withIdeButton', function() {
+      const source = `import React from 'react';
+
+const Button = () => {
+  return <button>Click me</button>;
+};
+
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(Button,");
+      expect(result).toMatch(/projectRoot: ['"][^'"]*['"]/); // Accept both single and double quotes and any valid path
+    });
+
+    it('should handle JSX files', function() {
+      mockContext.resourcePath = process.cwd() + '/src/components/Button.jsx';
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(Button,");
+    });
+
+    it('should handle components with existing imports', function() {
+      const source = `import React from 'react';
+import { useState } from 'react';
+
+const Counter = () => {
+  const [count, setCount] = useState(0);
+  return <div>{count}</div>;
+};
+
+export default Counter;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(Counter,");
+      
+      const lines = result.split('\n');
+      const importIndex = lines.findIndex(line => line.includes("import { withIdeButton }"));
+      const useStateIndex = lines.findIndex(line => line.includes("import { useState }"));
+      expect(importIndex).toBeGreaterThan(useStateIndex);
+    });
+
+    it('should handle components with no imports', function() {
+      const source = `const SimpleButton = () => {
+  return <button>Simple</button>;
+};
+
+export default SimpleButton;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(SimpleButton,");
+      expect(result.indexOf("import { withIdeButton }")).toBe(0);
+    });
+
+    it('should include relative path in the wrapper', function() {
+      mockContext.resourcePath = process.cwd() + '/src/components/ui/Button.tsx';
+      mockContext.getOptions.mockReturnValue({ 
+        enabled: true,
+        projectRoot: process.cwd()
+      });
+
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("'src/components/ui/Button.tsx'");
+    });
+
+    it('should use custom import path when specified', function() {
+      mockContext.getOptions.mockReturnValue({ 
+        enabled: true,
+        importPath: 'custom-ide-helper'
+      });
+
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'custom-ide-helper';");
+    });
+  });
+
+  describe('when disabled', function() {
+    it('should return source unchanged when enabled is false', function() {
+      mockContext.getOptions.mockReturnValue({ enabled: false });
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toBe(source);
+      expect(result).not.toContain('withIdeButton');
+    });
+
+    it('should return source unchanged when NODE_ENV is production and enabled not explicitly set', function() {
+      process.env.NODE_ENV = 'production';
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toBe(source);
+      expect(result).not.toContain('withIdeButton');
+    });
+  });
+
+  describe('file filtering', function() {
+    beforeEach(function() {
+      mockContext.getOptions.mockReturnValue({ enabled: true });
+    });
+
+    it('should not process files outside component directories', function() {
+      mockContext.resourcePath = process.cwd() + '/src/utils/helper.tsx';
+      const source = `const Helper = () => <div>Helper</div>;
+export default Helper;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toBe(source);
+      expect(result).not.toContain('withIdeButton');
+    });
+
+    it('should process files in custom component directories', function() {
+      mockContext.resourcePath = process.cwd() + '/src/widgets/MyWidget.tsx';
+      mockContext.getOptions.mockReturnValue({ 
+        enabled: true,
+        componentPaths: ['src/widgets'],
+        projectRoot: process.cwd()
+      });
+
+      const source = `const MyWidget = () => <div>Widget</div>;
+export default MyWidget;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain('withIdeButton');
+    });
+
+    it('should not process non-JSX/TSX files', function() {
+      mockContext.resourcePath = process.cwd() + '/src/components/Button.js';
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toBe(source);
+      expect(result).not.toContain('withIdeButton');
+    });
+  });
+
+  describe('component detection', function() {
+    beforeEach(function() {
+      mockContext.getOptions.mockReturnValue({ enabled: true });
+    });
+
+    it('should not process files without default export', function() {
+      const source = `export const Button = () => <button>Click</button>;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toBe(source);
+      expect(result).not.toContain('withIdeButton');
+    });
+
+    it('should not process files with lowercase default export', function() {
+      const source = `const button = () => <button>Click</button>;
+export default button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toBe(source);
+      expect(result).not.toContain('withIdeButton');
+    });
+
+    it('should process files with PascalCase default export', function() {
+      const source = `const MyButton = () => <button>Click</button>;
+export default MyButton;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain('withIdeButton');
+    });
+
+    it('should process direct export default function declarations', function() {
+      const source = `export default function MyCoolComponent() {
+  return <div>My Cool Component</div>;
+}`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(MyCoolComponent,");
+    });
+
+    it('should process direct export default arrow function expressions', function() {
+      const source = `export default () => {
+  return <div>Anonymous Component</div>;
+};`;
+
+      const result = loader.call(mockContext, source);
+
+      // Should process anonymous components using file path
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(");
+    });
+
+    it('should process anonymous function expressions', function() {
+      const source = `export default function() {
+  return <div>Anonymous Function Component</div>;
+}`;
+
+      const result = loader.call(mockContext, source);
+
+      // Should process anonymous function components
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(");
+    });
+
+    it('should process direct export default class components', function() {
+      const source = `import React from 'react';
+
+export default class MyClassComponent extends React.Component {
+  render() {
+    return <div>Class Component</div>;
+  }
+}`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(MyClassComponent,");
+    });
+
+    it('should not process lowercase direct export default functions', function() {
+      const source = `export default function myLowercaseComponent() {
+  return <div>Lowercase Component</div>;
+}`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toBe(source);
+      expect(result).not.toContain('withIdeButton');
+    });
+
+    it('should handle direct export default with complex function names', function() {
+      const source = `export default function MyComplexComponent123WithNumbers() {
+  return <div>Complex Named Component</div>;
+}`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(MyComplexComponent123WithNumbers,");
+    });
+
+    it('should handle function components with TypeScript', function() {
+      const source = `interface Props {
+  title: string;
+}
+
+export default function TypeScriptComponent(props: Props) {
+  return <div>{props.title}</div>;
+}`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(TypeScriptComponent,");
+      expect(result).toContain("function TypeScriptComponent(props: Props)");
+    });
+
+    it('should handle class components with TypeScript', function() {
+      const source = `import React, { Component } from 'react';
+
+interface Props {
+  name: string;
+}
+
+export default class TSClassComponent extends Component<Props> {
+  render() {
+    return <div>Hello {this.props.name}</div>;
+  }
+}`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("import { withIdeButton } from 'nextjs-ide-helper';");
+      expect(result).toContain("export default withIdeButton(TSClassComponent,");
+      expect(result).toContain("class TSClassComponent extends Component<Props>");
+    });
+  });
+
+  describe('already wrapped components', function() {
+    beforeEach(function() {
+      mockContext.getOptions.mockReturnValue({ enabled: true });
+    });
+
+    it('should not wrap components that are already wrapped', function() {
+      const source = `import { withIdeButton } from 'nextjs-ide-helper';
+
+const Button = () => <button>Click</button>;
+
+export default withIdeButton(Button, 'src/components/Button.tsx');`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toBe(source);
+      expect((result.match(/withIdeButton/g) || []).length).toBe(2);
+    });
+  });
+
+  describe('edge cases', function() {
+    beforeEach(function() {
+      mockContext.getOptions.mockReturnValue({ enabled: true });
+    });
+
+    it('should handle export default with semicolon', function() {
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("export default withIdeButton(Button,");
+      expect(result).not.toContain("export default Button;");
+    });
+
+    it('should handle export default without semicolon', function() {
+      const source = `const Button = () => <button>Click</button>
+export default Button`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("export default withIdeButton(Button,");
+      expect(result).not.toContain("export default Button");
+    });
+
+    it('should handle complex component names', function() {
+      const source = `const MyComplexComponent123 = () => <div>Complex</div>;
+export default MyComplexComponent123;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain("export default withIdeButton(MyComplexComponent123,");
+    });
+
+    it('should handle files with multiple component paths', function() {
+      mockContext.resourcePath = process.cwd() + '/lib/components/Button.tsx';
+      mockContext.getOptions.mockReturnValue({ 
+        enabled: true,
+        componentPaths: ['src/components', 'lib/components'],
+        projectRoot: process.cwd()
+      });
+
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain('withIdeButton');
+    });
+  });
+
+  describe('options handling', function() {
+    it('should use default options when none provided', function() {
+      mockContext.getOptions.mockReturnValue(null);
+      process.env.NODE_ENV = 'development';
+
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toContain('withIdeButton');
+      expect(result).toContain("'nextjs-ide-helper'");
+    });
+
+    it('should handle custom project root', function() {
+      const customRoot = '/custom/project';
+      mockContext.resourcePath = customRoot + '/src/components/Button.tsx';
+      mockContext.getOptions.mockReturnValue({ 
+        enabled: true,
+        projectRoot: customRoot
+      });
+
+      const source = `const Button = () => <button>Click</button>;
+export default Button;`;
+
+      const result = loader.call(mockContext, source);
+
+      expect(result).toMatch(new RegExp("projectRoot:\\s*['\"]" + customRoot.replace(/[.*+?^${}()|[\]\\]/g, '\\$&') + "['\"]"));
+      expect(result).toContain("'src/components/Button.tsx'");
+    });
+  });
+});

package/src/loader.js

@@ -1,4 +1,8 @@
 const path = require('path');
+const { parse } = require('@babel/parser');
+const traverse = require('@babel/traverse').default;
+const generate = require('@babel/generator').default;
+const t = require('@babel/types');
 
 /**
  * Webpack loader that automatically wraps React components with cursor buttons
@@ -30,50 +34,212 @@ module.exports = function cursorButtonLoader(source) {
     return source;
   }
 
-  // Check if it's already wrapped
+  // Check if it's already wrapped using simple string check for performance
   if (source.includes('withIdeButton')) {
     return source;
   }
 
-  // Check if it has a default export that looks like a React component
-  const defaultExportMatch = source.match(/export\s+default\s+(\w+)/);
-  if (!defaultExportMatch) {
+  let ast;
+  try {
+    // Parse the source code into an AST
+    ast = parse(source, {
+      sourceType: 'module',
+      plugins: [
+        'jsx',
+        'typescript',
+        'decorators-legacy',
+        'classProperties',
+        'objectRestSpread',
+        'asyncGenerators',
+        'functionBind',
+        'exportDefaultFrom',
+        'exportNamespaceFrom',
+        'dynamicImport',
+        'nullishCoalescingOperator',
+        'optionalChaining'
+      ]
+    });
+  } catch (error) {
+    // If parsing fails, return original source
+    return source;
+  }
+
+  let hasDefaultExport = false;
+  let defaultExportName = null;
+  let isAnonymousComponent = false;
+  let hasWithIdeButtonImport = false;
+  let lastImportPath = null;
+  let defaultExportPath = null;
+
+  // Traverse the AST to analyze the code
+  traverse(ast, {
+    ImportDeclaration(path) {
+      lastImportPath = path;
+      
+      // Check if withIdeButton is already imported
+      if (path.node.source.value === importPath) {
+        path.node.specifiers.forEach(spec => {
+          if (t.isImportSpecifier(spec) && spec.imported.name === 'withIdeButton') {
+            hasWithIdeButtonImport = true;
+          }
+        });
+      }
+    },
+    
+    ExportDefaultDeclaration(path) {
+      hasDefaultExport = true;
+      defaultExportPath = path;
+      
+      if (t.isIdentifier(path.node.declaration)) {
+        // export default ComponentName
+        defaultExportName = path.node.declaration.name;
+      } else if (t.isFunctionDeclaration(path.node.declaration)) {
+        if (path.node.declaration.id) {
+          // export default function ComponentName() {}
+          defaultExportName = path.node.declaration.id.name;
+        } else {
+          // export default function() {} - anonymous function
+          isAnonymousComponent = true;
+        }
+      } else if (t.isClassDeclaration(path.node.declaration) && path.node.declaration.id) {
+        // export default class ComponentName extends ... {}
+        defaultExportName = path.node.declaration.id.name;
+      } else if (t.isArrowFunctionExpression(path.node.declaration)) {
+        // export default () => {} - anonymous arrow function
+        isAnonymousComponent = true;
+      } else if (t.isVariableDeclaration(path.node.declaration)) {
+        // export default const ComponentName = ...
+        const declarator = path.node.declaration.declarations[0];
+        if (t.isIdentifier(declarator.id)) {
+          defaultExportName = declarator.id.name;
+        }
+      }
+      
+      // Stop traversal once we find the default export
+      path.stop();
+    }
+  });
+
+  // Check if we should process this file
+  if (!hasDefaultExport || (!defaultExportName && !isAnonymousComponent)) {
     return source;
   }
 
-  const componentName = defaultExportMatch[1];
-  
   // Check if component name starts with uppercase (React component convention)
-  if (!componentName || componentName[0] !== componentName[0].toUpperCase()) {
+  // Skip this check for anonymous components
+  if (defaultExportName && defaultExportName[0] !== defaultExportName[0].toUpperCase()) {
+    return source;
+  }
+
+  // If withIdeButton is already imported, don't process
+  if (hasWithIdeButtonImport) {
     return source;
   }
 
   // Get relative path from project root
   const relativePath = path.relative(projectRoot, filename);
 
-  // Find the last import to insert our import after it
-  const imports = source.match(/^import.*$/gm) || [];
-  
-  let modifiedSource;
-  
-  if (imports.length === 0) {
-    // No imports found, add at top
-    const importStatement = `import { withIdeButton } from '${importPath}';\n`;
-    const wrappedExport = `export default withIdeButton(${componentName}, '${relativePath}', { projectRoot: '${projectRoot}' });`;
-    modifiedSource = importStatement + source.replace(/export\s+default\s+\w+;?/, wrappedExport);
+  // Transform the AST
+  let modified = false;
+
+  // Add the withIdeButton import
+  const withIdeButtonImport = t.importDeclaration(
+    [t.importSpecifier(t.identifier('withIdeButton'), t.identifier('withIdeButton'))],
+    t.stringLiteral(importPath)
+  );
+
+  // Insert import after last existing import or at the beginning
+  if (lastImportPath) {
+    lastImportPath.insertAfter(withIdeButtonImport);
+    modified = true;
   } else {
-    const lastImportIndex = source.lastIndexOf(imports[imports.length - 1]) + imports[imports.length - 1].length;
-    
-    // Add import statement
-    const importStatement = `\nimport { withIdeButton } from '${importPath}';`;
-    
-    // Replace the export
-    const wrappedExport = `export default withIdeButton(${componentName}, '${relativePath}', { projectRoot: '${projectRoot}' });`;
+    // No imports found, add at the beginning
+    if (ast.program && ast.program.body && Array.isArray(ast.program.body)) {
+      ast.program.body.unshift(withIdeButtonImport);
+      modified = true;
+    }
+  }
+
+  // Replace the default export with wrapped version
+  if (defaultExportPath && modified) {
+    let wrappedCall;
     
-    modifiedSource = source.slice(0, lastImportIndex) + 
-                     importStatement + 
-                     source.slice(lastImportIndex).replace(/export\s+default\s+\w+;?/, wrappedExport);
+    if (isAnonymousComponent) {
+      // For anonymous components, wrap the original function/arrow function
+      let componentExpression;
+      
+      if (t.isFunctionDeclaration(defaultExportPath.node.declaration)) {
+        // Convert anonymous function declaration to function expression
+        componentExpression = t.functionExpression(
+          null, // id is null for anonymous
+          defaultExportPath.node.declaration.params,
+          defaultExportPath.node.declaration.body,
+          defaultExportPath.node.declaration.generator,
+          defaultExportPath.node.declaration.async
+        );
+      } else {
+        // For arrow functions and other expressions, use directly
+        componentExpression = defaultExportPath.node.declaration;
+      }
+      
+      wrappedCall = t.callExpression(
+        t.identifier('withIdeButton'),
+        [
+          componentExpression,
+          t.stringLiteral(relativePath),
+          t.objectExpression([
+            t.objectProperty(t.identifier('projectRoot'), t.stringLiteral(projectRoot))
+          ])
+        ]
+      );
+    } else {
+      // For named components, wrap with the component name identifier
+      wrappedCall = t.callExpression(
+        t.identifier('withIdeButton'),
+        [
+          t.identifier(defaultExportName),
+          t.stringLiteral(relativePath),
+          t.objectExpression([
+            t.objectProperty(t.identifier('projectRoot'), t.stringLiteral(projectRoot))
+          ])
+        ]
+      );
+    }
+
+    // If the export is a named function or class declaration, we need to handle it differently
+    if (!isAnonymousComponent && 
+        (t.isFunctionDeclaration(defaultExportPath.node.declaration) || 
+         t.isClassDeclaration(defaultExportPath.node.declaration))) {
+      // Insert the function/class declaration before the export
+      const declaration = defaultExportPath.node.declaration;
+      defaultExportPath.insertBefore(declaration);
+      
+      // Replace the export declaration with the wrapped call
+      defaultExportPath.node.declaration = wrappedCall;
+    } else {
+      // For other cases (identifier, variable declaration, anonymous functions), just replace the declaration
+      defaultExportPath.node.declaration = wrappedCall;
+    }
   }
 
-  return modifiedSource;
+  // Generate the transformed code
+  if (modified) {
+    try {
+      const result = generate(ast, {
+        retainLines: false,
+        compact: false,
+        jsescOption: {
+          quotes: 'single'
+        }
+      });
+      return result.code;
+    } catch (error) {
+      // If generation fails, return original source
+      return source;
+    }
+  }
+
+  // If we reach here, nothing was modified
+
+  return source;
 };