@codeyam/codeyam-cli 0.1.0-staging.8aea589 → 0.1.0-staging.bbe4da9

package/analyzer-template/project/writeScenarioComponents.ts

@@ -1,6 +1,7 @@
 import {
   Analysis,
   Entity,
+  EntityType,
   File,
   Scenario,
   Project,
@@ -25,6 +26,14 @@ import * as fs from 'fs';
 import * as path from 'path';
 import { LazyFileStore } from './LazyFileStore';
 
+/**
+ * Escape special regex characters in a string so it can be used as a literal pattern.
+ * This prevents characters like . * + ? ^ $ { } ( ) | [ ] \ from being interpreted as regex metacharacters.
+ */
+function escapeRegExp(str: string): string {
+  return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+}
+
 interface WriteScenarioComponentsArgs {
   project: Project;
   file: File;
@@ -43,14 +52,22 @@ interface WriteScenarioComponentsArgs {
   })[];
   writtenScenarioComponents?: { [key: string]: string[] };
   fileStore?: LazyFileStore;
+  /**
+   * When a default export needs to be accessible as a named export (e.g., when
+   * the parent imports it as `{ Name }` via an index re-export), pass the name here.
+   * A re-export statement `export { default as Name };` will be added to the output.
+   */
+  exportAsNamed?: string;
 }
 
 type EnhancedImportedExport = Omit<
   Entity['metadata']['importedExports'][0],
-  'entityType' | 'calls'
+  'calls' | 'entityType'
 > & {
   analysisId: string;
   isNodeModule?: boolean;
+  entityType?: EntityType; // Optional because nodeModuleImports don't have it
+  calls?: readonly string[]; // Optional - used for Zod schema detection heuristic
 };
 
 function hasMockStructure(
@@ -73,6 +90,7 @@ function hasMockStructure(
           importedExport.filePath
         ]?.[importedExport.name],
     )?.metadata?.mergedDataStructure?.dependencySchemas,
+    importedExport.entityType,
   );
 }
 
@@ -83,54 +101,340 @@ function isIndexPath(filePath: string) {
 }
 
 /**
- * Rewrite asset imports (CSS, images, fonts, etc.) to correct relative paths.
- * These imports aren't tracked as importedExports, so we need to handle them separately.
+ * Convert .d.ts type declaration content to stub implementations.
+ * .d.ts files only have type declarations which don't provide runtime exports.
+ * This function converts them to actual stub implementations.
  *
- * Handles patterns like:
+ * @param content - The .d.ts file content
+ * @param entityName - The name of the entity we're generating for (used to ensure it's exported)
+ * @returns The converted content with stub implementations
+ */
+function convertDtsToStubs(content: string, entityName: string): string {
+  let result = content;
+
+  // Handle object types FIRST (before simple types) since they span multiple lines
+  // and contain semicolons that would incorrectly terminate the simple type regex
+
+  // Convert "export declare const NAME: { ... };" (object type) to "export const NAME = {} as any;"
+  result = result.replace(
+    /export\s+declare\s+const\s+(\w+)\s*:\s*\{[^}]*\}\s*;/gs,
+    'export const $1 = {} as any;',
+  );
+
+  // Convert "export declare let NAME: { ... };" (object type) to "export let NAME = {} as any;"
+  result = result.replace(
+    /export\s+declare\s+let\s+(\w+)\s*:\s*\{[^}]*\}\s*;/gs,
+    'export let $1 = {} as any;',
+  );
+
+  // Convert "export declare var NAME: { ... };" (object type) to "export var NAME = {} as any;"
+  result = result.replace(
+    /export\s+declare\s+var\s+(\w+)\s*:\s*\{[^}]*\}\s*;/gs,
+    'export var $1 = {} as any;',
+  );
+
+  // Now handle simple types (non-object types)
+
+  // Convert "export declare const NAME: TYPE;" to "export const NAME = {} as any;"
+  result = result.replace(
+    /export\s+declare\s+const\s+(\w+)\s*:\s*[^;]+;/g,
+    'export const $1 = {} as any;',
+  );
+
+  // Convert "export declare let NAME: TYPE;" to "export let NAME = {} as any;"
+  result = result.replace(
+    /export\s+declare\s+let\s+(\w+)\s*:\s*[^;]+;/g,
+    'export let $1 = {} as any;',
+  );
+
+  // Convert "export declare var NAME: TYPE;" to "export var NAME = {} as any;"
+  result = result.replace(
+    /export\s+declare\s+var\s+(\w+)\s*:\s*[^;]+;/g,
+    'export var $1 = {} as any;',
+  );
+
+  // Convert "export declare function NAME(...): TYPE;" to "export function NAME() { return {} as any; }"
+  result = result.replace(
+    /export\s+declare\s+function\s+(\w+)\s*\([^)]*\)\s*:\s*[^;]+;/g,
+    'export function $1(...args: any[]) { return {} as any; }',
+  );
+
+  // Convert "export declare class NAME { ... }" to "export class NAME {}"
+  // This is tricky because class bodies can span multiple lines
+  result = result.replace(
+    /export\s+declare\s+class\s+(\w+)\s*(?:extends\s+[^{]+)?\{[^}]*\}/gs,
+    'export class $1 {}',
+  );
+
+  // Convert "declare const NAME: { ... }" (object type, non-exported) to "const NAME = {} as any;"
+  // Object types can span multiple lines and contain semicolons, so we handle them separately
+  // The 's' flag makes . match newlines, so this matches the entire object type block
+  result = result.replace(
+    /^declare\s+const\s+(\w+)\s*:\s*\{[^}]*\}\s*;/gms,
+    'const $1 = {} as any;',
+  );
+
+  // Convert "declare const NAME: TYPE;" (simple type, non-exported) to "const NAME = {} as any;"
+  // These are often used with "export { NAME }" at the end
+  // This must come AFTER the object type handler above
+  result = result.replace(
+    /^declare\s+const\s+(\w+)\s*:\s*[^;]+;/gm,
+    'const $1 = {} as any;',
+  );
+
+  // Convert "declare function NAME" (non-exported) to "function NAME() { return {} as any; }"
+  result = result.replace(
+    /^declare\s+function\s+(\w+)\s*\([^)]*\)\s*:\s*[^;]+;/gm,
+    'function $1(...args: any[]) { return {} as any; }',
+  );
+
+  // Convert "declare class NAME" (non-exported) to "class NAME {}"
+  result = result.replace(
+    /^declare\s+class\s+(\w+)\s*(?:extends\s+[^{]+)?\{[^}]*\}/gm,
+    'class $1 {}',
+  );
+
+  // Remove "export { }" at the end (empty export statement that marks module as having no exports)
+  result = result.replace(/export\s*\{\s*\}\s*;?\s*$/g, '');
+
+  // Keep export type and export interface statements as-is (they're valid in .ts)
+  // No transformation needed for these
+
+  console.log(
+    `CodeYam: Converted .d.ts content for entity "${entityName}". Result length: ${result.length}`,
+  );
+
+  return result;
+}
+
+/**
+ * Rewrite relative asset paths to correct locations when a file is moved.
+ *
+ * This function generically handles ANY string literal containing a relative path
+ * to an asset file, not just import statements. This catches patterns like:
  * - import "./globals.css"
  * - import styles from "./file.module.css"
- * - import logo from "./logo.png"
+ * - localFont({ src: "../fonts/font.woff2" })
+ * - Image src="../images/logo.png"
+ * - Any other string literal with a relative asset path
  */
 function rewriteAssetImports(
   fileContent: string,
   originalFilePath: string,
   newFilePath: string,
 ): string {
-  // Match imports with non-JS/TS extensions
+  // Asset file extensions that should have their paths rewritten
   const assetExtensions =
     'css|scss|sass|less|styl|png|jpg|jpeg|gif|svg|webp|ico|woff|woff2|ttf|eot|otf|mp4|webm|json';
-  const assetImportRegex = new RegExp(
-    `import\\s+(?:[\\w$]+\\s+from\\s+)?(['"])(\\.{1,2}/[^'"]+\\.(?:${assetExtensions}))\\1`,
+
+  // Match ANY string literal containing a relative path to an asset file
+  // This captures both single and double quoted strings that start with ./ or ../
+  const relativeAssetPathRegex = new RegExp(
+    `(['"])(\\.{1,2}/[^'"]*\\.(?:${assetExtensions}))\\1`,
     'g',
   );
 
   // Get the directory part of the original file path (relative to project root)
   const originalDirParts = originalFilePath.split('/').slice(0, -1);
 
-  return fileContent.replace(assetImportRegex, (match, quote, importPath) => {
-    // Resolve the asset path relative to the original file's directory
-    // We need to manually resolve the relative path without using path.resolve
-    // which would resolve against the actual filesystem
-    const importParts = importPath.split('/');
-    const resolvedParts = [...originalDirParts];
-
-    for (const part of importParts) {
-      if (part === '..') {
-        resolvedParts.pop();
-      } else if (part !== '.') {
-        resolvedParts.push(part);
+  return fileContent.replace(
+    relativeAssetPathRegex,
+    (match, quote, assetPath) => {
+      // Resolve the asset path relative to the original file's directory
+      // We need to manually resolve the relative path without using path.resolve
+      // which would resolve against the actual filesystem
+      const pathParts = assetPath.split('/');
+      const resolvedParts = [...originalDirParts];
+
+      for (const part of pathParts) {
+        if (part === '..') {
+          resolvedParts.pop();
+        } else if (part !== '.') {
+          resolvedParts.push(part);
+        }
       }
-    }
 
-    // Build the project-relative path to the asset
-    const absoluteAssetPath = resolvedParts.join('/');
+      // Build the project-relative path to the asset
+      const absoluteAssetPath = resolvedParts.join('/');
 
-    // Calculate the new relative path from new file location
-    const newRelativePath = getRelativePath(newFilePath, absoluteAssetPath);
+      // Calculate the new relative path from new file location
+      const newRelativePath = getRelativePath(newFilePath, absoluteAssetPath);
 
-    // Reconstruct the import statement
-    return match.replace(importPath, newRelativePath);
-  });
+      // Replace the path in the original match, preserving the quote style
+      return `${quote}${newRelativePath}${quote}`;
+    },
+  );
+}
+
+/**
+ * Rewrite relative TypeScript/JavaScript module imports when a file is moved.
+ *
+ * When a file is moved from one location to another (e.g., from [environmentId]/
+ * to _environmentId_/), relative imports like "./lib/organization" need to be
+ * rewritten to point to the correct location from the new path.
+ *
+ * This is similar to rewriteAssetImports but handles TypeScript/JavaScript modules
+ * instead of asset files (CSS, images, etc.).
+ */
+function rewriteRelativeModuleImports(
+  fileContent: string,
+  originalFilePath: string,
+  newFilePath: string,
+): string {
+  // Module file extensions that should have their relative paths rewritten
+  const moduleExtensions = 'ts|tsx|js|jsx|mjs|cjs';
+
+  // Match import statements with relative paths (starting with ./ or ../)
+  // This matches:
+  // - import { foo } from "./path"
+  // - import foo from "../path"
+  // - import * as foo from "./path"
+  // - import "./path" (side-effect imports)
+  // The path may or may not have a file extension
+  const relativeImportRegex = new RegExp(
+    `(from\\s+)(['"])(\\.\\.?\\/[^'"]+?)(?:\\.(?:${moduleExtensions}))?\\2`,
+    'g',
+  );
+
+  // Get the directory part of the original file path (relative to project root)
+  const originalDirParts = originalFilePath.split('/').slice(0, -1);
+
+  return fileContent.replace(
+    relativeImportRegex,
+    (match, fromKeyword, quote, importPath) => {
+      // Skip imports that already point to generated CodeYam files:
+      // 1. Scenario component files with SHA hashes (64 hex chars) and scenario name suffixes
+      //    e.g., "abc123def_EnvironmentLayout_Empty_Survey_No_Responses"
+      // 2. Mock data files in __codeyamMocks__ directory
+      //    e.g., "../__codeyamMocks__/MockData_Empty_Survey_No_Responses"
+      const scenarioFilePattern = /[a-f0-9]{64}_\w+_[A-Z][a-zA-Z_]+$/;
+      const mockDataPattern = /__codeyamMocks__\//;
+      if (
+        scenarioFilePattern.test(importPath) ||
+        mockDataPattern.test(importPath)
+      ) {
+        // This import already points to a generated file, don't rewrite it
+        return match;
+      }
+
+      // Resolve the import path relative to the original file's directory
+      const pathParts = importPath.split('/');
+      const resolvedParts = [...originalDirParts];
+
+      for (const part of pathParts) {
+        if (part === '..') {
+          resolvedParts.pop();
+        } else if (part !== '.') {
+          resolvedParts.push(part);
+        }
+      }
+
+      // Build the project-relative path to the module
+      const absoluteModulePath = resolvedParts.join('/');
+
+      // Calculate the new relative path from new file location
+      const newRelativePath = getRelativePath(newFilePath, absoluteModulePath);
+
+      // Replace the path in the original match, preserving the quote style
+      return `${fromKeyword}${quote}${newRelativePath}${quote}`;
+    },
+  );
+}
+
+/**
+ * Strip <html> and <body> tags from root layout files for Next.js App Router.
+ *
+ * When a root layout (app/layout.tsx) is generated as a scenario component and placed
+ * under the /static/ route, it becomes nested under the real app root layout. Having
+ * two layouts with <html> and <body> tags causes hydration errors:
+ *
+ *   "In HTML, <html> cannot be a child of <body>"
+ *   "You are mounting a new html component when a previous one has not first unmounted"
+ *
+ * This function removes the <html> and <body> wrapper tags while preserving the inner
+ * content, allowing the scenario layout to work correctly when nested.
+ *
+ * Before: <html lang="en"><body className={...}><main>{children}</main></body></html>
+ * After:  <div className={...}><main>{children}</main></div>
+ */
+function stripHtmlBodyTags(
+  fileContent: string,
+  filePath: string,
+  framework: ProjectFramework,
+): string {
+  // Only apply to Next.js App Router root layouts
+  if (framework !== ProjectFramework.Next) {
+    return fileContent;
+  }
+
+  // Check if this is a root layout file (typically app/layout.tsx or apps/*/app/layout.tsx)
+  // Root layouts are at the app directory level, not in subdirectories like (app)/ or routes/
+  const pathParts = filePath.split('/');
+  const fileName = pathParts[pathParts.length - 1];
+
+  // Must be a layout file
+  if (!fileName?.startsWith('layout.')) {
+    return fileContent;
+  }
+
+  // Check if it's at the app root level (parent directory is 'app')
+  const parentDir = pathParts[pathParts.length - 2];
+  if (parentDir !== 'app') {
+    return fileContent;
+  }
+
+  // Check if this file actually contains <html> and <body> tags
+  if (!/<html[^>]*>/.test(fileContent) || !/<body[^>]*>/.test(fileContent)) {
+    return fileContent;
+  }
+
+  console.log(
+    `CodeYam: Stripping <html> and <body> tags from root layout: ${filePath}`,
+  );
+
+  // Extract the body className/attributes if any, to preserve styling
+  const bodyMatch = fileContent.match(/<body([^>]*)>/);
+  const bodyAttributes = bodyMatch?.[1]?.trim() || '';
+
+  // Replace the JSX structure:
+  // <html ...><body ...>CONTENT</body></html>
+  // With: <div ...>CONTENT</div>
+  //
+  // This regex handles:
+  // 1. Opening <html> tag with any attributes
+  // 2. Opening <body> tag with any attributes (captured for the wrapper div)
+  // 3. Content between body tags (captured)
+  // 4. Closing </body> and </html> tags
+  const htmlBodyRegex =
+    /(<html[^>]*>\s*<body)([^>]*)(>)([\s\S]*?)(<\/body>\s*<\/html>)/g;
+
+  let result = fileContent.replace(
+    htmlBodyRegex,
+    (match, _htmlBody, bodyAttrs, closeBracket, content, _closing) => {
+      // Create a wrapper div with the body's attributes
+      const attrs = bodyAttrs?.trim() || '';
+      if (attrs) {
+        return `(<div${bodyAttrs}${closeBracket}${content}</div>)`;
+      } else {
+        return `(<>${content}</>)`;
+      }
+    },
+  );
+
+  // If the regex didn't match (perhaps due to different formatting),
+  // try a more lenient approach - just remove the tags
+  if (result === fileContent) {
+    // Remove <html ...> opening tag
+    result = result.replace(/<html[^>]*>\s*/g, '');
+    // Remove </html> closing tag
+    result = result.replace(/\s*<\/html>/g, '');
+    // Replace <body ...> with <div ...> to preserve attributes
+    result = result.replace(/<body([^>]*)>/g, '<div$1>');
+    // Replace </body> with </div>
+    result = result.replace(/<\/body>/g, '</div>');
+  }
+
+  return result;
 }
 
 function addMockToContent(
@@ -157,7 +461,11 @@ function addMockToContent(
       rootAnalysis.metadata?.mergedDataStructure?.dependencySchemas;
   }
 
-  const mockCode = constructMockCode(importedExport.name, dependencySchemas);
+  const mockCode = constructMockCode(
+    importedExport.name,
+    dependencySchemas,
+    importedExport.entityType,
+  );
 
   if (!mockCode) {
     console.log(
@@ -204,8 +512,13 @@ function addMockToContent(
       );
       fileContent = fileContent.replace(entireImportRegExp, '');
     } else {
+      // Escape regex special characters in importPath (e.g., brackets in [environmentId])
+      const escapedImportPath = importPath.replace(
+        /[.*+?^${}()|[\]\\]/g,
+        '\\$&',
+      );
       const importRegExp = new RegExp(
-        `(import(?:(?!${firstPart}|from|import)[\\s\\S])*?)${firstPart}(?:,\\s*)?((?:(?!import)[\\s\\S])*?from\\s+['"]${importPath}['"])`,
+        `(import(?:(?!${firstPart}|from|import)[\\s\\S])*?)${firstPart}(?:,\\s*)?((?:(?!import)[\\s\\S])*?from\\s+['"]${escapedImportPath}['"])`,
         'm',
       );
 
@@ -225,6 +538,38 @@ function addMockToContent(
         '',
       );
     }
+  } else {
+    // Fallback: When importPath is undefined (e.g., path alias not in mapping),
+    // we still need to remove the import to avoid "name defined multiple times" errors.
+    // Search for the entity name in any import statement and remove it.
+    const importedExportNameParts = splitOutsideParenthesesAndArrays(
+      importedExport.name,
+    );
+    const firstPart = importedExportNameParts[0];
+
+    // Match the entity name in any named import and remove just that name
+    // This handles imports like: import { useEnvironment, otherThing } from "any/path"
+    // Removes useEnvironment but keeps otherThing
+    const namedImportRegExp = new RegExp(
+      `(import\\s*\\{[^}]*?)\\b${escapeRegExp(firstPart)}\\b(?:,\\s*|\\s*,)?((?:[^}]*\\}\\s*from\\s*['"][^'"]*['"];?))`,
+      'm',
+    );
+
+    if (importedExportNameParts.length > 1) {
+      // Rename the import instead of removing (for destructured access patterns)
+      fileContent = fileContent.replace(
+        namedImportRegExp,
+        `$1${firstPart}__cyOriginal$2`,
+      );
+    } else {
+      fileContent = fileContent.replace(namedImportRegExp, '$1$2');
+    }
+
+    // Clean up empty imports that might result from removing the only import
+    fileContent = fileContent.replace(
+      /import\s*\{\s*\}\s*from\s+['"][^'"]*['"];?\s*\n?/g,
+      '',
+    );
   }
 
   if (fileContent.indexOf('import { scenarios } from') === -1) {
@@ -409,6 +754,7 @@ export default async function writeScenarioComponents({
   namespaceMocks,
   writtenScenarioComponents = {},
   fileStore,
+  exportAsNamed,
 }: WriteScenarioComponentsArgs): Promise<{
   scenarioComponentPaths: string[];
   writtenScenarioComponents: { [key: string]: string[] };
@@ -529,6 +875,29 @@ export default async function writeScenarioComponents({
   );
 
   let fileContent = file.content;
+
+  // Handle .d.ts files: convert type declarations to stub implementations
+  // .d.ts files only have type declarations (e.g., "export declare const logger")
+  // which don't provide runtime exports. We need to generate actual stub implementations.
+  if (file.path.endsWith('.d.ts')) {
+    console.log(
+      `CodeYam: Converting .d.ts file to stub implementations: ${file.path}`,
+    );
+    fileContent = convertDtsToStubs(fileContent, entity.name);
+  }
+
+  // Extract "use client" or "use server" directive FIRST, before any modifications
+  // This ensures the directive isn't buried by prepended imports
+  const directiveMatch = fileContent.match(
+    /^(\s*["']use (client|server)["'];?\s*\n?)/,
+  );
+  let extractedDirective: string | null = null;
+  if (directiveMatch) {
+    extractedDirective = directiveMatch[1].trim();
+    // Remove the directive from fileContent - we'll add it back at the very end
+    fileContent = fileContent.slice(directiveMatch[0].length);
+  }
+
   const fileAnalyzer = projectAnalyzer.getFileAnalyzer(file);
   const importMapping = fileAnalyzer.getRelativeImportMappings();
 
@@ -643,10 +1012,12 @@ export default async function writeScenarioComponents({
     (importedExportEntity as Entity).localFilePath = importedExport.filePath;
 
     if (!importedExport.isMocked) {
+      // Use resolvedFilePath (actual entity location) for scenario component writing.
+      // This is critical for re-exports: when importing from index.tsx which re-exports
+      // from a nested file, we need to write scenario components for the actual file
+      // where the entity code lives, not the re-export file.
       const importedExportFilePath =
-        importedExportEntity.localFilePath === file.path
-          ? (importedExport.resolvedFilePath ?? importedExport.filePath)
-          : importedExport.filePath;
+        importedExport.resolvedFilePath ?? importedExport.filePath;
 
       if (importedExportFilePath !== file.path) {
         if (
@@ -668,17 +1039,41 @@ export default async function writeScenarioComponents({
             ? fileStore.getByPath(importedExportFilePath)
             : project.files.find((f) => f.path === importedExportFilePath);
 
+          // Fallback: if file at resolvedFilePath doesn't exist, try filePath
+          // This handles cases where the actual entity file isn't in project.files
+          // (e.g., type re-exports where only the index file is tracked)
+          if (
+            !fileNotMocked &&
+            importedExport.resolvedFilePath &&
+            importedExport.filePath !== importedExport.resolvedFilePath
+          ) {
+            const fallbackPath = importedExport.filePath;
+            fileNotMocked = fileStore
+              ? fileStore.getByPath(fallbackPath)
+              : project.files.find((f) => f.path === fallbackPath);
+          }
+
+          // Skip if file still not found - can't write scenario component without the file
+          if (!fileNotMocked) {
+            continue;
+          }
+
           // Ensure content is loaded for the file
           if (
             fileNotMocked &&
             fileStore &&
-            !fileStore.isContentLoaded(importedExportFilePath)
+            !fileStore.isContentLoaded(fileNotMocked.path)
           ) {
-            fileNotMocked = await fileStore.ensureContent(
-              importedExportFilePath,
-            );
+            fileNotMocked = await fileStore.ensureContent(fileNotMocked.path);
           }
 
+          // When a default export is imported as named (via index re-export), we need
+          // to add a named re-export to the scenario component so the import works.
+          // e.g., if file has `export default X` but parent does `import { X } from '...'`
+          const needsNamedReExport =
+            importedExport.resolvedIsDefault === true &&
+            importedExport.isDefault === false;
+
           const {
             scenarioComponentPaths: newScenarioComponentPaths,
             writtenScenarioComponents: updatedWrittenScenarioComponents,
@@ -696,6 +1091,8 @@ export default async function writeScenarioComponents({
             namespaceMocks,
             writtenScenarioComponents,
             fileStore,
+            // Pass the import name so we can add `export { default as Name };`
+            exportAsNamed: needsNamedReExport ? importedExport.name : undefined,
           });
           writtenScenarioComponents = updatedWrittenScenarioComponents;
           scenarioComponentPaths.push(...newScenarioComponentPaths);
@@ -706,21 +1103,67 @@ export default async function writeScenarioComponents({
       importedExport.name !== entity.name
     ) {
       // For same-file dependencies:
-      // - Strip if isMocked is true (e.g., Remix/Next loaders, server-side code)
-      // - Strip if we have a mock structure we can generate
-      // - Only add mock code if we have a mock structure
-      const shouldStrip =
-        importedExport.isMocked ||
-        hasMockStructure(importedExport, fileAnalyses);
-
-      if (shouldStrip) {
+      // - Strip and replace if we have a mock structure we can generate
+      // - Strip and stub if isMocked is true AND it's a callable entity (visual, library, functionCall)
+      // - Preserve as-is for data entities (like Zod schemas) that need their methods
+      const hasMock = hasMockStructure(importedExport, fileAnalyses);
+      const entityType = importedExport.entityType;
+      // Data and type entities should be preserved - they're not callable and may have methods
+      // that stubbing would break (e.g., Zod schemas with .superRefine())
+      const isDataEntity = entityType === 'data' || entityType === 'type';
+
+      // Heuristic: Zod schemas are often misclassified as 'library' but should be preserved
+      // Detect by: name starts with Z + uppercase letter, AND has Zod method calls
+      const looksLikeZodSchema =
+        entityType === 'library' &&
+        /^Z[A-Z]/.test(importedExport.name) &&
+        importedExport.calls?.some((call: string) =>
+          /\.(superRefine|refine|transform|default|optional|nullable|array|object|string|number|boolean|parse|safeParse)\s*\(/.test(
+            call,
+          ),
+        );
+
+      if (looksLikeZodSchema) {
+        console.log(
+          `CodeYam: Detected Zod schema "${importedExport.name}" (misclassified as library) - will preserve`,
+        );
+      }
+
+      // Callable entities can be safely stubbed (but not Zod schemas)
+      const isCallable =
+        !isDataEntity && !looksLikeZodSchema && entityType !== undefined;
+
+      // Determine what action to take
+      const shouldStripAndReplace = hasMock;
+      const shouldStripAndStub =
+        !hasMock && importedExport.isMocked && isCallable;
+      const shouldPreserve = !hasMock && importedExport.isMocked && !isCallable;
+
+      // Log warning if entityType is undefined for a mocked same-file dependency
+      // This could indicate an analysis issue where entityType wasn't properly set
+      if (importedExport.isMocked && entityType === undefined) {
+        console.warn(
+          `CodeYam: WARNING - Same-file dependency "${importedExport.name}" is isMocked but has no entityType. ` +
+            `Treating as non-callable (will preserve). File: ${file.path}`,
+        );
+      }
+
+      if (shouldPreserve) {
+        // For data entities (like Zod schemas), don't strip or stub - preserve the original
+        // This ensures schema methods like .superRefine() continue to work
+        console.log(
+          `CodeYam: Preserving ${importedExport.name} (entityType: ${entityType}) - not stripping data entities`,
+        );
+        // Don't modify fileContent - keep the original code
+      } else if (shouldStripAndReplace || shouldStripAndStub) {
+        // Strip the original code
         const entityCode = fileAnalyzer.getEntityCode(importedExport.name);
         if (entityCode) {
           fileContent = fileContent.replace(entityCode, '');
         }
 
-        // Only add mock if we have a mock structure to generate
-        if (hasMockStructure(importedExport, fileAnalyses)) {
+        if (shouldStripAndReplace) {
+          // Add mock from mock structure
           fileContent = addMockToContent(
             fileContent,
             importedExport,
@@ -729,18 +1172,40 @@ export default async function writeScenarioComponents({
             relativeMocksDir,
             scenario.name,
           );
+        } else if (shouldStripAndStub) {
+          // Generate a simple stub mock that returns scenario data.
+          // This prevents ReferenceError at runtime when the stripped
+          // function is called (e.g., local helper functions like getInitialProps).
+          const functionName = importedExport.name;
+          console.log(
+            `CodeYam: Generating stub mock for ${functionName} (entityType: ${entityType}) in ${file.path}`,
+          );
+
+          // Add scenarios import if not present
+          if (fileContent.indexOf('import { scenarios } from') === -1) {
+            const mockDataPath = `${relativeMocksDir}/MockData_${safeFileName(scenario.name)}`;
+            const importStatement = `import { scenarios } from "${mockDataPath}";`;
+            fileContent = `${fileContent}\n\n\n${importStatement}`;
+          }
+
+          // Generate a simple stub that returns the scenario data for this function
+          const stubMock = `\n\n// Stub mock for local function without mock structure\nfunction ${functionName}(...args) {\n  return scenarios().data()?.["${functionName}()"];\n}`;
+          fileContent += stubMock;
         }
       }
     }
 
     // Try to find the import mapping using different key formats
     // The import mapping may use either absolute or relative paths as keys
-    // IMPORTANT: Use filePath (where it's imported FROM), not resolvedFilePath (where entity lives)
-    // For re-exports like ~codeyam/types, filePath="packages/types/index.ts",
-    // resolvedFilePath="packages/types/src/types/Commit.ts"
-    // We need to generate files at the import location (index.ts), not the entity location
+    // Use filePath (import source) for looking up the import mapping
+    // because that's how imports are written in source code
     const mockFilePathRelative = importedExport.filePath;
 
+    // For scenario component paths, prefer resolvedFilePath (where entity actually lives)
+    // because that's where scenario component files are written
+    const scenarioFilePathRelative =
+      importedExport.resolvedFilePath ?? importedExport.filePath;
+
     // Try looking up with different key formats (absolute, relative, resolved absolute)
     // Filter out undefined keys to avoid misleading undefined lookups
     const lookupKeys = [
@@ -794,10 +1259,28 @@ export default async function writeScenarioComponents({
     } else {
       // For non-mocked imports, rewrite the import path to point to the generated scenario file
       // Use fileStore for O(1) lookup when available
-      const fileNotMocked = fileStore
-        ? fileStore.getByPath(mockFilePathRelative)
-        : project.files.find((f) => f.path === mockFilePathRelative);
-      const fileName = mockFilePathRelative.split('/').pop();
+      // First try resolvedFilePath (where entity lives), fallback to filePath (import source)
+      let fileNotMocked = fileStore
+        ? fileStore.getByPath(scenarioFilePathRelative)
+        : project.files.find((f) => f.path === scenarioFilePathRelative);
+
+      // Track which path we're actually using for the scenario component
+      let actualScenarioFilePathRelative = scenarioFilePathRelative;
+      if (
+        !fileNotMocked &&
+        importedExport.resolvedFilePath &&
+        importedExport.filePath !== importedExport.resolvedFilePath
+      ) {
+        // Fallback to filePath if file at resolvedFilePath doesn't exist
+        actualScenarioFilePathRelative = importedExport.filePath;
+        fileNotMocked = fileStore
+          ? fileStore.getByPath(actualScenarioFilePathRelative)
+          : project.files.find(
+              (f) => f.path === actualScenarioFilePathRelative,
+            );
+      }
+
+      const fileName = actualScenarioFilePathRelative.split('/').pop();
       const fileNotMockedIsIndex = isIndexPath(fileNotMocked?.path);
       const mockFilePath = isFrameworkRoute(
         fileNotMocked,
@@ -810,7 +1293,7 @@ export default async function writeScenarioComponents({
             rootFile,
             entity: importedExportEntity,
             rootAnalysis,
-            scenarioComponentPath: mockFilePathRelative,
+            scenarioComponentPath: actualScenarioFilePathRelative,
             project,
             framework,
             scenario,
@@ -818,7 +1301,7 @@ export default async function writeScenarioComponents({
             .split('.')
             .slice(0, -1)
             .join('.')
-        : mockFilePathRelative.replace(
+        : actualScenarioFilePathRelative.replace(
             `${fileName}`,
             `${importedExportEntity.sha}_${fileNotMockedIsIndex ? 'index_' : ''}${safeFileName(importedExportEntity.name)}_${safeFileName(scenario.name)}`,
           );
@@ -829,11 +1312,45 @@ export default async function writeScenarioComponents({
       // Otherwise, skip rewriting (we can't find the import statement without knowing what to search for)
       if (mockImportMapping) {
         const importRegExp = new RegExp(
-          `${mockImportMapping.replace(/([$()])/g, '\\$1')}(['"])`,
+          `${escapeRegExp(mockImportMapping)}(['"])`,
           'g',
         );
 
-        fileContent = fileContent.replace(importRegExp, `${path}$1`);
+        // Check if the original import path still exists in the file
+        // If it was already rewritten by another entity from the same module,
+        // we need to add a separate import statement for this entity
+        if (fileContent.match(importRegExp)) {
+          fileContent = fileContent.replace(importRegExp, `${path}$1`);
+        } else {
+          // The import path was already rewritten - add a new import for this entity
+          // This handles cases where multiple entities are imported from the same index file
+          // (e.g., import { A, B, C } from '@pkg') and each has its own scenario file
+          const entityImportName = importedExport.name;
+          const newImport = `import { ${entityImportName} } from '${path}';`;
+
+          // First, try to remove this entity from the already-rewritten grouped import
+          // This prevents duplicate/conflicting imports
+          // Match patterns like "EntityName," or ", EntityName" or "EntityName" (if only one)
+          // Note: entityImportName needs escaping since JS identifiers can contain $ (a regex metacharacter)
+          const escapedEntityName = escapeRegExp(entityImportName);
+          const removeFromGroupedImportPatterns = [
+            new RegExp(`\\b${escapedEntityName}\\s*,\\s*`, 'g'), // "EntityName, "
+            new RegExp(`\\s*,\\s*${escapedEntityName}\\b`, 'g'), // ", EntityName"
+          ];
+          for (const pattern of removeFromGroupedImportPatterns) {
+            fileContent = fileContent.replace(pattern, '');
+          }
+
+          // Add the new import at the beginning of fileContent
+          // Note: The header comment (// Scenario:) doesn't exist yet - it's prepended at writeFile time
+          // So prepending here puts the import right after the header in the final output
+          //
+          // IMPORTANT: We cannot use indexOf('// Scenario:') + '\n\n' here because:
+          // 1. The header doesn't exist during processing
+          // 2. indexOf would return -1, and indexOf('\n\n', -1) starts from 0
+          // 3. This could find a \n\n inside a function body, inserting the import there!
+          fileContent = newImport + '\n' + fileContent;
+        }
       }
     }
   }
@@ -852,6 +1369,36 @@ export default async function writeScenarioComponents({
     );
   }
 
+  // Rewrite node_module imports that have universal mocks
+  // Universal mocks create mock files at __codeyamMocks__/{safeFileName}.tsx
+  // We need to rewrite imports like `import { logger } from "@formbricks/logger"`
+  // to `import { logger } from "../__codeyamMocks__/_formbricks_logger"`
+  const universalMocks = project.metadata?.universalMocks ?? [];
+  const nodeModuleUniversalMocks = universalMocks.filter(
+    (mock) => mock.nodeModule && mock.content,
+  );
+
+  for (const universalMock of nodeModuleUniversalMocks) {
+    const originalPath = universalMock.filePath;
+    // Create the mock file name using the same safeFileName function as writeUniversalMocks
+    const safeMockFileName = safeFileName(originalPath);
+    const mockFileRelativePath = `${relativeMocksDir}/${safeMockFileName}`;
+
+    // Match both single and double quotes, and handle both named and default imports
+    // Pattern 1: import { x, y } from "module"
+    // Pattern 2: import x from "module"
+    // Pattern 3: import * as x from "module"
+    const importRegex = new RegExp(
+      `(import\\s+(?:\\{[^}]*\\}|\\*\\s+as\\s+\\w+|\\w+)\\s+from\\s+)['"]${originalPath.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}['"]`,
+      'g',
+    );
+
+    fileContent = fileContent.replace(
+      importRegex,
+      `$1'${mockFileRelativePath}'`,
+    );
+  }
+
   if (
     rootAnalysis.entitySha === entity.sha &&
     entity.metadata?.notExported &&
@@ -864,6 +1411,17 @@ export default async function writeScenarioComponents({
     }
   }
 
+  // When a default export is imported as named via re-export (e.g., index.tsx re-exports
+  // a default export as named), we need to add a named export so the parent's
+  // import statement works: `import { Name } from './path'`
+  // The local variable name is the same as exportAsNamed (e.g., ScenarioEditor)
+  if (exportAsNamed) {
+    const namedExport = `export { ${exportAsNamed} };`;
+    if (!fileContent.includes(namedExport)) {
+      fileContent = `${fileContent}\n\n${namedExport}`;
+    }
+  }
+
   const basePath = safeFolder(file.path.split('/').slice(0, -1).join('/'));
   const extension = file.name.split('.').pop();
   // Include scenario name in path to allow multiple scenarios to coexist
@@ -882,6 +1440,10 @@ export default async function writeScenarioComponents({
     });
   }
 
+  // Strip <html> and <body> tags from root layout files for Next.js
+  // These tags cause hydration errors when the scenario layout is nested under the real root
+  fileContent = stripHtmlBodyTags(fileContent, file.path, framework);
+
   // Rewrite asset imports (CSS, images, fonts, etc.) to correct relative paths
   // The original file path is relative to PROJECT_RELATIVE_PATH, the new path is scenarioComponentPath
   fileContent = rewriteAssetImports(
@@ -890,6 +1452,15 @@ export default async function writeScenarioComponents({
     scenarioComponentPath,
   );
 
+  // Rewrite relative TypeScript/JavaScript module imports to correct relative paths
+  // This handles cases where the file is moved (e.g., from [environmentId]/ to _environmentId_/)
+  // and relative imports like "./lib/organization" need to be rewritten
+  fileContent = rewriteRelativeModuleImports(
+    fileContent,
+    `${PROJECT_RELATIVE_PATH}/${file.path}`,
+    scenarioComponentPath,
+  );
+
   console.log(
     'Writing scenario component',
     file.path,
@@ -909,10 +1480,19 @@ export default async function writeScenarioComponents({
 // Scenario: ${scenario.id} - ${scenario.name}
 `;
 
-  await writeFile(
-    scenarioComponentPath,
-    `${scenarioComponentComment}\n\n${fileContent}`,
-  );
+  // Use the directive that was extracted at the beginning of processing
+  // This ensures it stays at the very top even after imports are prepended
+  let finalContent: string;
+  if (extractedDirective) {
+    finalContent = `${extractedDirective}\n\n${scenarioComponentComment}\n\n${fileContent}`;
+    console.log(
+      `CodeYam: Placed "${extractedDirective}" directive at top of file: ${scenarioComponentPath}`,
+    );
+  } else {
+    finalContent = `${scenarioComponentComment}\n\n${fileContent}`;
+  }
+
+  await writeFile(scenarioComponentPath, finalContent);
   scenarioComponentPaths.push(scenarioComponentPath);
 
   console.log('CodeYam [writeScenarioComponents]: Generated scenario files', {