@codeyam/codeyam-cli 0.1.0-staging.da5baf5 → 0.1.0-staging.e38f7bd

package/background/src/lib/virtualized/project/writeScenarioComponents.js

@@ -1,16 +1,74 @@
+import { ProjectFramework, } from "../../../../../packages/types/index.js";
 import writeFile from "../common/writeFile.js";
 import { isFrameworkRoute, getFrameworkRoutePath, safeFileName, PROJECT_RELATIVE_PATH, } from "../../../../../packages/utils/index.js";
 import { getRelativePath, safeFolder } from "../../../../../packages/generate/index.js";
 import constructMockCode from "./constructMockCode.js";
 import { splitOutsideParenthesesAndArrays } from "../../../../../packages/ai/index.js";
-import { loadEntities } from "../../../../../packages/supabase/index.js";
 import * as fs from 'fs';
 import * as path from 'path';
+import ts from 'typescript';
+/**
+ * Find the end position of the last import/export-from statement using TypeScript AST.
+ * This is more reliable than regex for handling multiline imports, comments, etc.
+ *
+ * @param content - The file content to analyze
+ * @returns The position after the last import statement, or 0 if no imports found
+ */
+function findEndOfImports(content) {
+    try {
+        const sourceFile = ts.createSourceFile('temp.ts', content, ts.ScriptTarget.Latest, true);
+        let lastImportEnd = 0;
+        // Visit all top-level statements to find import/export declarations
+        for (const statement of sourceFile.statements) {
+            if (ts.isImportDeclaration(statement)) {
+                // Get the end position including any trailing semicolon/newline
+                lastImportEnd = statement.end;
+            }
+            else if (ts.isExportDeclaration(statement) &&
+                statement.moduleSpecifier) {
+                // Handle "export { x } from 'module'" style re-exports
+                lastImportEnd = statement.end;
+            }
+        }
+        return lastImportEnd;
+    }
+    catch (error) {
+        // If parsing fails, fall back to 0 (will append at end)
+        console.warn('CodeYam: Failed to parse file for import detection:', error);
+        return 0;
+    }
+}
+/**
+ * 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) {
+    return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+}
+/**
+ * Map nested dist paths to src paths.
+ * Some build tools create nested structures like:
+ *   packages/X/dist/X/src/Y.d.ts
+ * when the source is at:
+ *   packages/X/src/Y.ts
+ * This function normalizes such paths to the src location.
+ */
+function mapNestedDistToSrc(filePath) {
+    // Pattern: packages/X/dist/X/src/Y -> packages/X/src/Y
+    const nestedDistMatch = filePath.match(/^(packages\/[^/]+)\/dist\/[^/]+\/src\/(.+)$/);
+    if (nestedDistMatch) {
+        const [, packagePath, relativePath] = nestedDistMatch;
+        return `${packagePath}/src/${relativePath}`
+            .replace(/\.d\.ts$/, '.ts')
+            .replace(/\.js$/, '.ts');
+    }
+    return filePath;
+}
 function hasMockStructure(importedExport, fileAnalyses) {
     if (!fileAnalyses?.some((a) => a.entity.metadata?.isolatedDataStructure?.dependencySchemas)) {
         return false;
     }
-    return !!constructMockCode(importedExport.name, fileAnalyses.find((a) => !!a.metadata?.mergedDataStructure?.dependencySchemas?.[importedExport.filePath]?.[importedExport.name])?.metadata?.mergedDataStructure?.dependencySchemas);
+    return !!constructMockCode(importedExport.name, fileAnalyses.find((a) => !!a.metadata?.mergedDataStructure?.dependencySchemas?.[importedExport.filePath]?.[importedExport.name])?.metadata?.mergedDataStructure?.dependencySchemas, importedExport.entityType);
 }
 function isIndexPath(filePath) {
     if (!filePath)
@@ -18,6 +76,56 @@ function isIndexPath(filePath) {
     const fileName = filePath.split('/').pop();
     return fileName.startsWith('index.');
 }
+/**
+ * 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.
+ *
+ * @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, entityName) {
+    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.
  *
@@ -59,7 +167,142 @@ function rewriteAssetImports(fileContent, originalFilePath, newFilePath) {
         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, originalFilePath, newFilePath) {
+    // 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.js"
+    // - import foo from "../path.js"
+    // - import * as foo from "./path/index.js"
+    // - 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"
+        //    e.g., "abc123def_EnvironmentLayout_Survey_not_found_404" (scenario names can contain digits)
+        // 2. Mock data files in __codeyamMocks__ directory
+        //    e.g., "../__codeyamMocks__/MockData_Empty_Survey_No_Responses"
+        const scenarioFilePattern = /[a-f0-9]{64}_\w+_[A-Z]\w*$/;
+        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, filePath, framework) {
+    // 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;
+}
+// Version for tracking deployments - increment when making changes
+const WRITE_SCENARIO_COMPONENTS_VERSION = '2.1.0-ast-based-import-detection';
 function addMockToContent(fileContent, importedExport, fileAnalyses, rootAnalysis, relativeMocksDir, scenarioName, importPath) {
+    console.log(`CodeYam [v${WRITE_SCENARIO_COMPONENTS_VERSION}]: Adding mock for ${importedExport.name} from ${importedExport.filePath}`);
     // First try to find dependency schemas in fileAnalyses (for same-file dependencies)
     let dependencySchemas = fileAnalyses.find((a) => !!a.metadata?.mergedDataStructure?.dependencySchemas?.[importedExport.filePath]?.[importedExport.name])?.metadata?.mergedDataStructure?.dependencySchemas;
     // If not found in fileAnalyses, use root analysis's dependency schemas
@@ -68,11 +311,118 @@ function addMockToContent(fileContent, importedExport, fileAnalyses, rootAnalysi
         dependencySchemas =
             rootAnalysis.metadata?.mergedDataStructure?.dependencySchemas;
     }
-    const mockCode = constructMockCode(importedExport.name, dependencySchemas);
-    if (!mockCode) {
+    // Check if we have multiple calls with different variable names
+    // This requires generating separate mock functions for each call site
+    const hasMultipleCallsWithVariables = importedExport.calls &&
+        importedExport.calls.length > 1 &&
+        importedExport.callVariableNames &&
+        importedExport.callVariableNames.length === importedExport.calls.length;
+    // DEBUG: Log import info for multiple calls debugging
+    if (importedExport.name === 'useFetcher' ||
+        importedExport.name?.includes('Fetcher')) {
+        console.log('CodeYam DEBUG: addMockToContent useFetcher import:', JSON.stringify({
+            name: importedExport.name,
+            calls: importedExport.calls,
+            callVariableNames: importedExport.callVariableNames,
+            hasMultipleCallsWithVariables,
+            isMocked: importedExport.isMocked,
+        }, null, 2));
+    }
+    let mockCode;
+    const variableMockCodes = [];
+    if (hasMultipleCallsWithVariables) {
+        // Generate separate mock functions for each variable-qualified call
+        for (let i = 0; i < importedExport.calls.length; i++) {
+            const variableName = importedExport.callVariableNames[i];
+            if (!variableName)
+                continue;
+            // Generate mock code for this specific call using variable-qualified name
+            // Format: "variableName <- functionName" (reads as "variableName receives from functionName")
+            const qualifiedName = `${variableName} <- ${importedExport.name}`;
+            const variableMockCode = constructMockCode(qualifiedName, dependencySchemas, importedExport.entityType);
+            if (variableMockCode) {
+                variableMockCodes.push(variableMockCode);
+                // Replace the call site with the variable-specific mock function
+                // e.g., useFetcher<BranchEntityDiffResult>() -> useFetcher_entityDiffFetcher()
+                // e.g., useFetcher() -> useFetcher_reportFetcher()
+                const callSignature = importedExport.calls[i];
+                // Escape special regex characters in the call signature
+                const escapedCallSignature = callSignature.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+                // Create regex that matches the call (with optional whitespace variations)
+                const callRegex = new RegExp(escapedCallSignature.replace(/\s+/g, '\\s*'), 'g');
+                const mockFunctionName = `${importedExport.name}_${variableName}`;
+                fileContent = fileContent.replace(callRegex, `${mockFunctionName}()`);
+            }
+        }
+    }
+    else {
+        // Check if there's a single variable name that receives this function's return value
+        // Note: callVariableNames tracks the ACTUAL variable assignments (e.g., "const loadError = useLoaderData()")
+        // while calls can contain many entries including method chain accesses (e.g., "useLoaderData().config...")
+        // We only care about how many distinct variables receive the return value.
+        const singleCallVariableName = importedExport.callVariableNames &&
+            importedExport.callVariableNames.length === 1
+            ? importedExport.callVariableNames[0]
+            : undefined;
+        if (singleCallVariableName) {
+            // For single variable assignments, use the variable-qualified key for data lookup
+            // but keep the original function name (no need for unique function names when there's only one assignment)
+            const qualifiedKey = `${singleCallVariableName} <- ${importedExport.name}`;
+            mockCode = `const ${importedExport.name}ReturnValue = scenarios().data()?.["${qualifiedKey}"];
+
+function ${importedExport.name}() {
+  return ${importedExport.name}ReturnValue;
+}`;
+        }
+        else {
+            // Check if any analysis (fileAnalyses or rootAnalysis) has this function's data
+            // under a variable-qualified key. The entity that CALLS the function (e.g., FileTableRow)
+            // has the dataForMocks with the variable-qualified key, not the root analysis (e.g., GitView).
+            let variableQualifiedKey;
+            // First check fileAnalyses (the analyses for the entity being written)
+            for (const analysis of fileAnalyses) {
+                const dataForMocks = analysis.metadata?.scenariosDataStructure?.dataForMocks;
+                if (dataForMocks) {
+                    variableQualifiedKey = Object.keys(dataForMocks).find((key) => {
+                        const match = key.match(/^([a-zA-Z_][a-zA-Z0-9_]*)\s*<-\s*(.+)$/);
+                        return match && match[2] === importedExport.name;
+                    });
+                    if (variableQualifiedKey) {
+                        break;
+                    }
+                }
+            }
+            // If not found in fileAnalyses, fall back to rootAnalysis
+            if (!variableQualifiedKey) {
+                const dataForMocks = rootAnalysis.metadata?.scenariosDataStructure?.dataForMocks;
+                if (dataForMocks) {
+                    variableQualifiedKey = Object.keys(dataForMocks).find((key) => {
+                        const match = key.match(/^([a-zA-Z_][a-zA-Z0-9_]*)\s*<-\s*(.+)$/);
+                        return match && match[2] === importedExport.name;
+                    });
+                }
+            }
+            if (variableQualifiedKey) {
+                // Use the variable-qualified key found in the analysis
+                mockCode = `const ${importedExport.name}ReturnValue = scenarios().data()?.["${variableQualifiedKey}"];
+
+function ${importedExport.name}() {
+  return ${importedExport.name}ReturnValue;
+}`;
+            }
+            else {
+                // Original behavior for calls without variable names
+                mockCode = constructMockCode(importedExport.name, dependencySchemas, importedExport.entityType);
+            }
+        }
+    }
+    // Combine all mock codes
+    const allMockCodes = variableMockCodes.length > 0 ? variableMockCodes.join('\n\n') : mockCode;
+    if (!allMockCodes) {
         console.log('CodeYam Error: Mock code not found', JSON.stringify({
             importedExportFilePath: importedExport.filePath,
             importedExportEntityName: importedExport.name,
+            hasMultipleCallsWithVariables,
             analysisIds: fileAnalyses.map((a) => ({
                 id: a.id,
                 filePath: a.filePath,
@@ -93,7 +443,9 @@ function addMockToContent(fileContent, importedExport, fileAnalyses, rootAnalysi
             fileContent = fileContent.replace(entireImportRegExp, '');
         }
         else {
-            const importRegExp = new RegExp(`(import(?:(?!${firstPart}|from|import)[\\s\\S])*?)${firstPart}(?:,\\s*)?((?:(?!import)[\\s\\S])*?from\\s+['"]${importPath}['"])`, 'm');
+            // 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+['"]${escapedImportPath}['"])`, 'm');
             if (importedExportNameParts.length > 1) {
                 fileContent = fileContent.replace(importRegExp, `$1${firstPart}__cyOriginal$2`);
             }
@@ -105,13 +457,51 @@ function addMockToContent(fileContent, importedExport, fileAnalyses, rootAnalysi
             fileContent = fileContent.replace(/import\s*\{\s*\}\s*from\s+['"][^'"]*['"];?\s*\n?/g, '');
         }
     }
+    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, '');
+    }
+    // Find the last import position using TypeScript AST (more reliable than regex).
+    // This ensures we insert both the scenarios import AND the mock after the original imports,
+    // not at the end of the file.
+    const lastImportEnd = findEndOfImports(fileContent);
+    // Build the content to insert: scenarios import (if needed) + mock
+    let insertContent = '';
     if (fileContent.indexOf('import { scenarios } from') === -1) {
         // Use scenario-specific MockData file to allow multiple scenarios to coexist
         const mockDataPath = `${relativeMocksDir}/MockData_${safeFileName(scenarioName)}`;
-        const importStatement = `import { scenarios } from "${mockDataPath}";`;
-        fileContent = `${fileContent}\n\n\n${importStatement}`;
+        insertContent += `\n\nimport { scenarios } from "${mockDataPath}";`;
+    }
+    insertContent += `\n\n// Mock constructed using mocksDataStructure from analyses: ${fileAnalyses.map((a) => a.id).join(', ')}\n${allMockCodes}`;
+    if (lastImportEnd > 0) {
+        // Insert after the last original import
+        console.log(`CodeYam [v${WRITE_SCENARIO_COMPONENTS_VERSION}]: Inserting mock at position ${lastImportEnd} (after imports), not at end`);
+        fileContent =
+            fileContent.slice(0, lastImportEnd) +
+                insertContent +
+                fileContent.slice(lastImportEnd);
+    }
+    else {
+        // Fallback: append at end if no imports found
+        console.log(`CodeYam [v${WRITE_SCENARIO_COMPONENTS_VERSION}]: No imports found, appending mock at end`);
+        fileContent += insertContent;
     }
-    fileContent += `\n\n// Mock constructed using mocksDataStructure from analyses: ${fileAnalyses.map((a) => a.id).join(', ')}\n${mockCode}`;
     return fileContent;
 }
 // Debug utility to capture complex arguments for testing
@@ -324,6 +714,48 @@ export default async function writeScenarioComponents({ project, file, entity, r
     const nodeModuleImports = allImportedExports.filter((nodeModuleImport, index, self) => nodeModuleImport.isNodeModule &&
         index === self.findIndex((nmi) => nmi.name === nodeModuleImport.name));
     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);
+        // After basic .d.ts conversion, enhance the entity's stub with proper mock code
+        // if we have dependency schema data showing its methods/properties.
+        // This prevents "X.method is not a function" errors at runtime.
+        //
+        // NOTE: constructMockCode iterates through ALL file paths in dependencySchemas
+        // to find the entity, so we don't need to check for an exact file path match.
+        // The entity might be stored under a different path (e.g., import alias like
+        // @formbricks/logger vs the actual file path packages/logger/dist/index.d.ts).
+        const dependencySchemas = rootAnalysis.metadata?.mergedDataStructure?.dependencySchemas;
+        if (dependencySchemas) {
+            const mockCode = constructMockCode(entity.name, dependencySchemas, entity.entityType);
+            if (mockCode) {
+                // Replace the simple `export const entityName = {} as any;` with the full mock code
+                // The mockCode generates `const entityName = { ... }` so we need to add `export`
+                const stubPattern = new RegExp(`export\\s+(const|let|var)\\s+${escapeRegExp(entity.name)}\\s*=\\s*\\{\\s*\\}\\s+as\\s+any;?`);
+                if (fileContent.match(stubPattern)) {
+                    // Add export keyword to the mock code since constructMockCode generates `const name = ...`
+                    const exportMockCode = mockCode.replace(/^(const|let|var)\s+/, 'export $1 ');
+                    fileContent = fileContent.replace(stubPattern, exportMockCode);
+                    // Add scenarios import at the top for the mock to work
+                    const relativeMocksDir = getRelativePath(file.path, mocksDir);
+                    const mockDataPath = `${relativeMocksDir}/MockData_${safeFileName(scenario.name)}`;
+                    fileContent = `import { scenarios } from "${mockDataPath}";\n\n${fileContent}`;
+                }
+            }
+        }
+    }
+    // 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 = 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();
     let filePath = file.path;
@@ -371,37 +803,41 @@ export default async function writeScenarioComponents({ project, file, entity, r
         const resolvedFilePathForLookup = importedExport.resolvedFilePath
             ? normalizePath(`${projectAnalyzer.dirPath}/${importedExport.resolvedFilePath}`)
             : undefined;
-        let importedExportEntity = context.getEntity(importedExport.resolvedFilePath ?? importedExport.filePath, importedExport.resolvedName ?? importedExport.name) ?? context.getEntity(importedExport.resolvedFilePath, 'default');
-        if (!importedExportEntity) {
-            console.log('CodeYam Error: Entity not found', JSON.stringify({
-                filePath: file.path,
-                entityName: entity.name,
-                analysisIds: fileAnalyses.map((a) => a.id),
-                analysisFilePath: fileAnalyses[0]?.filePath,
-                analysisEntityName: fileAnalyses.map((a) => a.entityName),
-                importedExport,
-                contextEntitiesForFilePath: context.getEntityNamesForFile(importedExport.filePath),
-                contextEntitiesForResolvedFilePath: context.getEntityNamesForFile(importedExport.resolvedFilePath),
-            }, null, 2));
-            // Why is this necessary? The entity should be in the context.
-            const candidateEntities = await loadEntities({
-                projectId: project.id,
-                filePaths: [importedExport.resolvedFilePath ?? importedExport.filePath],
-                names: [importedExport.resolvedName ?? importedExport.name],
-            });
-            importedExportEntity = candidateEntities.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime())[0];
-            if (!importedExportEntity) {
-                throw new Error('importedExportEntity not found in writeScenarioComponents');
+        const importedExportEntity = context.getEntity(importedExport.resolvedFilePath ?? importedExport.filePath, importedExport.resolvedName ?? importedExport.name) ?? context.getEntity(importedExport.resolvedFilePath, 'default');
+        // Check if we actually need the entity for this import
+        const importedExportFilePath = importedExport.resolvedFilePath ?? importedExport.filePath;
+        const isSameFile = importedExportFilePath === file.path;
+        // We only need the entity for non-mocked, different-file imports:
+        // - Same-file imports: skip to the isCurrentFile check later (line ~1203)
+        // - Mocked imports: use importedExport.entityType, not importedExportEntity
+        // - Non-mocked different-file: need entity for type check, recursive call, and path rewriting
+        if (!importedExportEntity && !isSameFile && !importedExport.isMocked) {
+            // Check if the file exists in project.files - if not, skip this import.
+            // This happens when a file is excluded from project.files (e.g., Prisma-generated client.ts)
+            // but entities still have importedExports referencing it.
+            const fileExistsInProject = fileStore
+                ? fileStore.getByPath(importedExportFilePath)
+                : project.files.find((f) => f.path === importedExportFilePath);
+            if (!fileExistsInProject) {
+                // File not in project.files - likely auto-generated or excluded. Skip it.
+                continue;
             }
+            // All entities should be pre-loaded by loadReadyToBeCaptured's gatherAllReferencedEntities.
+            // If we're here, that function has a bug or an edge case was missed.
+            throw new Error(`Entity not found in context: ${importedExport.resolvedFilePath ?? importedExport.filePath}:${importedExport.resolvedName ?? importedExport.name}\n` +
+                `Looking up for: ${file.path}:${entity.name}\n` +
+                `Context entities for file: ${JSON.stringify(context.getEntityNamesForFile(importedExport.filePath))}\n` +
+                `Context entities for resolved: ${JSON.stringify(context.getEntityNamesForFile(importedExport.resolvedFilePath))}`);
+        }
+        if (importedExportEntity) {
+            importedExportEntity.localFilePath = importedExport.filePath;
         }
-        importedExportEntity.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 = importedExport.resolvedFilePath ?? importedExport.filePath;
-            if (importedExportFilePath !== file.path) {
+            if (!isSameFile) {
                 if (!writtenScenarioComponents[importedExportFilePath]?.includes(importedExport.name)) {
                     // Skip recursion for type/data entities - they don't need scenario components
                     // Only recurse for visual and library entities that need mocking
@@ -464,35 +900,94 @@ export default async function writeScenarioComponents({ project, file, entity, r
         else if (importedExport.filePath === file.path &&
             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) => /\.(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) {
+                    if (entityCode) {
+                        fileContent = fileContent.replace(entityCode, '');
+                    }
+                    // Add mock from mock structure
                     fileContent = addMockToContent(fileContent, importedExport, fileAnalyses, rootAnalysis, relativeMocksDir, scenario.name);
                 }
-                else if (importedExport.isMocked) {
-                    // No mock structure available, but isMocked is true.
+                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;
+                    // Check if the original entity was exported - if so, the stub should be exported too
+                    // This ensures other files that import this entity can still access it
+                    const wasExported = entityCode?.trimStart().startsWith('export');
+                    const exportKeyword = wasExported ? 'export ' : '';
+                    // Generate a chainable stub mock using a Proxy
+                    // This handles cases where the entity is used with method chains like:
+                    //   actionClient.use(...).use(...)
+                    // The Proxy returns itself for any property access, making it infinitely chainable.
+                    // When called as a function, it returns the scenario data.
+                    const stubMock = `// Stub mock for local function without mock structure
+// Uses a Proxy to support method chaining (e.g., ${functionName}.use(...))
+const _${functionName}Handler = {
+  get: function(target, prop) {
+    if (prop === 'then') return undefined; // Prevent Promise-like behavior
+    return (...args) => new Proxy(() => scenarios().data()?.["${functionName}()"], _${functionName}Handler);
+  },
+  apply: function(target, thisArg, args) {
+    return scenarios().data()?.["${functionName}()"];
+  }
+};
+${exportKeyword}const ${functionName} = new Proxy(() => scenarios().data()?.["${functionName}()"], _${functionName}Handler);`;
+                    // Replace the original entity code with the stub mock IN PLACE
+                    // This is critical to maintain the original ordering of declarations.
+                    // If we append the stub at the end, code that uses this entity before
+                    // the stub definition will fail with "Cannot access X before initialization".
+                    if (entityCode) {
+                        fileContent = fileContent.replace(entityCode, stubMock);
+                    }
+                    else {
+                        // Fallback: append at end if we can't find the original code
+                        fileContent += `\n\n${stubMock}`;
+                    }
                 }
             }
         }
@@ -502,8 +997,9 @@ export default async function writeScenarioComponents({ project, file, entity, r
         // 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;
+        // because that's where scenario component files are written.
+        // Map nested dist paths to src paths so import paths point to actual source files.
+        const scenarioFilePathRelative = mapNestedDistToSrc(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 = [
@@ -582,12 +1078,26 @@ export default async function writeScenarioComponents({ project, file, entity, r
             // If we have an import mapping, use it to find the import string to replace
             // 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')}(['"])`, 'g');
+                const importRegExp = new RegExp(`${escapeRegExp(mockImportMapping)}(['"])`, 'g');
+                // Check if there are mocked entities from the same file
+                // If so, we can't replace the entire import path - the mocked entities will be inlined
+                // and we need to keep the original import path for this non-mocked entity
+                const sameFileImports = sortedImportedExports.filter((ie) => ie.filePath === importedExport.filePath);
+                const hasMockedSameFileImports = sameFileImports.some((ie) => ie.isMocked && ie.name !== importedExport.name);
                 // 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`);
+                    if (hasMockedSameFileImports) {
+                        // There are mocked entities from the same file that will be inlined
+                        // Don't rewrite the import path - the original import is still needed for this non-mocked entity
+                        // The mocked entities will be inlined, so the original import statement will remain
+                        // with only the non-mocked entities after the mocked ones are removed
+                    }
+                    else {
+                        // No mocked entities from the same file - safe to replace entire path
+                        fileContent = fileContent.replace(importRegExp, `${path}$1`);
+                    }
                 }
                 else {
                     // The import path was already rewritten - add a new import for this entity
@@ -598,26 +1108,24 @@ export default async function writeScenarioComponents({ project, file, entity, r
                     // 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${entityImportName}\\s*,\\s*`, 'g'), // "EntityName, "
-                        new RegExp(`\\s*,\\s*${entityImportName}\\b`, 'g'), // ", EntityName"
+                        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 top of the file (after the auto-generated comment)
-                    const commentEndIndex = fileContent.indexOf('\n\n', fileContent.indexOf('// Scenario:'));
-                    if (commentEndIndex !== -1) {
-                        fileContent =
-                            fileContent.slice(0, commentEndIndex + 2) +
-                                newImport +
-                                '\n' +
-                                fileContent.slice(commentEndIndex + 2);
-                    }
-                    else {
-                        // Fallback: add at the very beginning
-                        fileContent = newImport + '\n' + fileContent;
-                    }
+                    // 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;
                 }
             }
         }
@@ -627,6 +1135,24 @@ export default async function writeScenarioComponents({ project, file, entity, r
             continue;
         fileContent = addMockToContent(fileContent, nodeModuleImport, fileAnalyses, rootAnalysis, relativeMocksDir, scenario.name, importMapping[nodeModuleImport.filePath] ?? nodeModuleImport.filePath);
     }
+    // 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.js"`
+    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 &&
         entity.name !== 'default') {
@@ -662,9 +1188,16 @@ export default async function writeScenarioComponents({ project, file, entity, r
             scenario,
         });
     }
+    // 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(fileContent, `${PROJECT_RELATIVE_PATH}/${file.path}`, 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, entity.name, scenarioComponentPath, fileContent.length);
     const scenarioComponentComment = `// This file is auto-generated by CodeYam. Do not edit this file manually.
 // This file contains content for a scenario component:
@@ -676,7 +1209,17 @@ export default async function writeScenarioComponents({ project, file, entity, r
 // Entity: ${rootAnalysis.entitySha} ${rootAnalysis.entityName}
 // 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;
+    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', {
         entityName: entity.name,