diff --git a/apps/web/src/app/components/get-imported-components-for.tsx b/apps/web/src/app/components/get-imported-components-for.tsx
index 4153b987db..073cc44cf7 100644
--- a/apps/web/src/app/components/get-imported-components-for.tsx
+++ b/apps/web/src/app/components/get-imported-components-for.tsx
@@ -1,8 +1,11 @@
 import { promises as fs } from 'node:fs';
 import path from 'node:path';
+import { spec } from 'node:test/reporters';
 import { parse } from '@babel/parser';
 import traverse from '@babel/traverse';
+import * as allReactEmailComponents from '@react-email/components';
 import { render } from '@react-email/components';
+import * as allReactResponsiveComponents from '@responsive-email/react-email';
 import { z } from 'zod';
 import { Layout } from '../../../components/_components/layout';
 import type { Category, Component } from '../../../components/structure';
@@ -26,7 +29,10 @@ const ComponentModule = z.object({
   component: z.record(z.string(), z.any()),
 });
 
-const getComponentCodeFrom = (fileContent: string): string => {
+const getComponentCodeFrom = (
+  componentName: string,
+  fileContent: string,
+): string => {
   const parsedContents = parse(fileContent, {
     sourceType: 'unambiguous',
     strictMode: false,
@@ -35,6 +41,12 @@ const getComponentCodeFrom = (fileContent: string): string => {
   });
 
   let componentCode: string | undefined;
+
+  const nativeComponents = Object.keys(allReactEmailComponents);
+  const usedNativeComponents = new Set<string>();
+  const responsiveEmailComponents = Object.keys(allReactResponsiveComponents);
+  const usedResponsiveEmailComponents = new Set<string>();
+
   traverse(parsedContents, {
     VariableDeclarator({ node }) {
       if (
@@ -48,16 +60,57 @@ const getComponentCodeFrom = (fileContent: string): string => {
         }
       }
     },
+    ImportSpecifier({ node }) {
+      const specifier =
+        node.imported.type === 'Identifier'
+          ? node.imported.name
+          : node.imported.value;
+      if (nativeComponents.includes(specifier)) {
+        usedNativeComponents.add(specifier);
+      }
+      if (responsiveEmailComponents.includes(specifier)) {
+        usedResponsiveEmailComponents.add(specifier);
+      }
+    },
   });
 
   if (!componentCode) {
     throw new Error('Could not find the source code for the component');
   }
 
-  return componentCode
-    .split(/\r\n|\r|\n/)
-    .map((line) => line.replace(/^\s{2}/, ''))
-    .join('\n');
+  componentCode = `export default function ${componentName}() {
+  return ${componentCode}
+}`;
+
+  let importStatements = 'import * as React from "react";\n';
+
+  function generateImportStatement(specifiers: string[], source: string) {
+    let statement = `import { ${specifiers.join(', ')} } from "${source}";\n`;
+
+    if (statement.length > 80) {
+      statement = `import {${specifiers.map((specifier) => `\n  ${specifier}`).join(',')}\n} from "${source}";\n`;
+    }
+
+    return statement;
+  }
+
+  if (usedNativeComponents.size > 0) {
+    importStatements += generateImportStatement(
+      Array.from(usedNativeComponents),
+      '@react-email/components',
+    );
+  }
+
+  if (usedResponsiveEmailComponents.size > 0) {
+    importStatements += generateImportStatement(
+      Array.from(responsiveEmailComponents),
+      '@responsive-email/react-email',
+    );
+  }
+
+  componentCode = `${importStatements}\n${componentCode}`;
+
+  return componentCode;
 };
 
 export const getComponentElement = async (
@@ -81,6 +134,11 @@ export const getImportedComponent = async (
   const dirpath = getComponentPathFromSlug(component.slug);
   const variantFilenames = await fs.readdir(dirpath);
 
+  const componentName = component.slug.replaceAll(
+    /(?:^|-)([a-z])/g,
+    (_match, letter: string) => letter.toUpperCase(),
+  );
+
   if (variantFilenames.length === 1 && variantFilenames[0] === 'index.tsx') {
     const filePath = path.join(dirpath, 'index.tsx');
     const element = <Layout>{await getComponentElement(filePath)}</Layout>;
@@ -88,7 +146,7 @@ export const getImportedComponent = async (
       pretty: true,
     });
     const fileContent = await fs.readFile(filePath, 'utf8');
-    const code = getComponentCodeFrom(fileContent);
+    const code = getComponentCodeFrom(componentName, fileContent);
     return {
       ...component,
       code: {
@@ -116,7 +174,10 @@ export const getImportedComponent = async (
 
   variantFilenames.forEach((variantFilename, index) => {
     const variantKey = variantFilename.replace('.tsx', '') as CodeVariant;
-    codePerVariant[variantKey] = getComponentCodeFrom(fileContents[index]);
+    codePerVariant[variantKey] = getComponentCodeFrom(
+      componentName,
+      fileContents[index],
+    );
   });
 
   const element = <Layout>{elements[0]}</Layout>;
diff --git a/apps/web/src/components/code-block.tsx b/apps/web/src/components/code-block.tsx
index df0a4f7330..ed2ccb5802 100644
--- a/apps/web/src/components/code-block.tsx
+++ b/apps/web/src/components/code-block.tsx
@@ -1,7 +1,6 @@
 import classNames from 'classnames';
 import type { Language } from 'prism-react-renderer';
 import { Highlight } from 'prism-react-renderer';
-import * as React from 'react';
 
 const theme = {
   plain: {
@@ -63,19 +62,19 @@ export const CodeBlock: React.FC<Readonly<CodeBlockProps>> = ({
 
           <pre className="p-4 font-mono">
             {tokens.map((line, i) => {
-              const lineProps = getLineProps({ line, key: i });
+              const lineProps = getLineProps({ line });
 
               return (
                 <div
-                  key={i}
                   {...lineProps}
+                  key={i}
                   className={classNames('whitespace-pre', {
                     "before:mr-2 before:text-slate-11 before:content-['$']":
                       language === 'bash' && tokens.length === 1,
                   })}
                 >
                   {line.map((token, key) => {
-                    const tokenProps = getTokenProps({ token, key });
+                    const tokenProps = getTokenProps({ token });
 
                     const isException =
                       token.content === 'from' &&
@@ -84,11 +83,7 @@ export const CodeBlock: React.FC<Readonly<CodeBlockProps>> = ({
                       ? [...token.types, 'key-white']
                       : token.types;
 
-                    return (
-                      <React.Fragment key={key}>
-                        <span {...tokenProps} />
-                      </React.Fragment>
-                    );
+                    return <span {...tokenProps} key={key} />;
                   })}
                 </div>
               );
diff --git a/apps/web/src/components/component-code-view.tsx b/apps/web/src/components/component-code-view.tsx
index 51948ef179..10a82b71cb 100644
--- a/apps/web/src/components/component-code-view.tsx
+++ b/apps/web/src/components/component-code-view.tsx
@@ -6,8 +6,6 @@ import { useStoredState } from '@/hooks/use-stored-state';
 import { convertUrisIntoUrls } from '@/utils/convert-uris-into-urls';
 import * as Select from '@radix-ui/react-select';
 import * as Tabs from '@radix-ui/react-tabs';
-import * as allReactEmailComponents from '@react-email/components';
-import * as allReactResponsiveComponents from '@responsive-email/react-email';
 import {
   CheckIcon,
   ChevronDownIcon,
@@ -45,34 +43,6 @@ export const ComponentCodeView = ({
   }
   code = convertUrisIntoUrls(code);
 
-  if (selectedLanguage === 'react') {
-    const importsReactResponsive = extractReactComponents(
-      code,
-      Object.keys(allReactResponsiveComponents),
-    );
-
-    const importsReactEmail = extractReactComponents(
-      code,
-      Object.keys(allReactEmailComponents),
-    );
-
-    let importStatements = '';
-
-    if (importsReactEmail.length > 0) {
-      importStatements += `import { ${importsReactEmail.join(
-        ', ',
-      )} } from "@react-email/components";\n`;
-    }
-
-    if (importsReactResponsive.length > 0) {
-      importStatements += `import { ${importsReactResponsive.join(
-        ', ',
-      )} } from "@responsive-email/react-email";\n`;
-    }
-
-    code = `${importStatements}\n${code}`;
-  }
-
   const onCopy = () => {
     void navigator.clipboard.writeText(code);
     setIsCopied(true);
@@ -211,25 +181,3 @@ const ReactVariantSelect = ({
     </Select.Root>
   );
 };
-
-/**
- * Extracts React component names from a string of React/JSX code
- */
-const extractReactComponents = (
-  code: string,
-  supportedComponents: string[],
-): string[] => {
-  const componentPattern =
-    /(?:<|import\s+\{?\s*)(?<componentName>[A-Z][a-zA-Z0-9]*)/g;
-  const matches = Array.from(code.matchAll(componentPattern));
-
-  const componentNames = Array.from(
-    new Set(
-      matches
-        .map((match) => match.groups?.componentName)
-        .filter((name): name is string => Boolean(name)),
-    ),
-  );
-
-  return componentNames.filter((name) => supportedComponents.includes(name));
-};