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

package/analyzer-template/project/writeScenarioComponents.ts

@@ -21,11 +21,51 @@ import {
 import { getRelativePath, safeFolder } from '~codeyam/generate';
 import constructMockCode from './constructMockCode';
 import { splitOutsideParenthesesAndArrays } from '~codeyam/ai';
-import { loadEntities } from '~codeyam/supabase';
 import * as fs from 'fs';
 import * as path from 'path';
+import ts from 'typescript';
 import { LazyFileStore } from './LazyFileStore';
 
+/**
+ * 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: string): number {
+  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.
@@ -34,6 +74,28 @@ function escapeRegExp(str: string): string {
   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: string): string {
+  // 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;
+}
+
 interface WriteScenarioComponentsArgs {
   project: Project;
   file: File;
@@ -62,12 +124,13 @@ interface WriteScenarioComponentsArgs {
 
 type EnhancedImportedExport = Omit<
   Entity['metadata']['importedExports'][0],
-  'calls' | 'entityType'
+  'calls' | 'entityType' | 'callVariableNames'
 > & {
   analysisId: string;
   isNodeModule?: boolean;
   entityType?: EntityType; // Optional because nodeModuleImports don't have it
   calls?: readonly string[]; // Optional - used for Zod schema detection heuristic
+  callVariableNames?: readonly string[]; // Variable names receiving the return value of each call
 };
 
 function hasMockStructure(
@@ -305,9 +368,10 @@ function rewriteRelativeModuleImports(
       // 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][a-zA-Z_]+$/;
+      const scenarioFilePattern = /[a-f0-9]{64}_\w+_[A-Z]\w*$/;
       const mockDataPattern = /__codeyamMocks__\//;
       if (
         scenarioFilePattern.test(importPath) ||
@@ -437,6 +501,9 @@ function stripHtmlBodyTags(
   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: string,
   importedExport: EnhancedImportedExport,
@@ -446,6 +513,9 @@ function addMockToContent(
   scenarioName: string,
   importPath?: string,
 ) {
+  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) =>
@@ -461,19 +531,157 @@ function addMockToContent(
       rootAnalysis.metadata?.mergedDataStructure?.dependencySchemas;
   }
 
-  const mockCode = constructMockCode(
-    importedExport.name,
-    dependencySchemas,
-    importedExport.entityType,
-  );
+  // 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;
 
-  if (!mockCode) {
+  // 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: string | undefined;
+  const variableMockCodes: string[] = [];
+
+  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: string | undefined;
+
+      // 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,
@@ -572,14 +780,38 @@ function addMockToContent(
     );
   }
 
+  // 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}";`;
   }
 
-  fileContent += `\n\n// Mock constructed using mocksDataStructure from analyses: ${fileAnalyses.map((a) => a.id).join(', ')}\n${mockCode}`;
+  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;
+  }
 
   return fileContent;
 }
@@ -884,6 +1116,45 @@ export default async function writeScenarioComponents({
       `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
@@ -961,65 +1232,54 @@ export default async function writeScenarioComponents({
         )
       : undefined;
 
-    let importedExportEntity =
+    const 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',
-        );
+    // 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))}`,
+      );
     }
 
-    (importedExportEntity as Entity).localFilePath = importedExport.filePath;
+    if (importedExportEntity) {
+      (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 =
-        importedExport.resolvedFilePath ?? importedExport.filePath;
-
-      if (importedExportFilePath !== file.path) {
+      if (!isSameFile) {
         if (
           !writtenScenarioComponents[importedExportFilePath]?.includes(
             importedExport.name,
@@ -1158,11 +1418,11 @@ export default async function writeScenarioComponents({
       } else if (shouldStripAndReplace || shouldStripAndStub) {
         // Strip the original code
         const entityCode = fileAnalyzer.getEntityCode(importedExport.name);
-        if (entityCode) {
-          fileContent = fileContent.replace(entityCode, '');
-        }
 
         if (shouldStripAndReplace) {
+          if (entityCode) {
+            fileContent = fileContent.replace(entityCode, '');
+          }
           // Add mock from mock structure
           fileContent = addMockToContent(
             fileContent,
@@ -1188,9 +1448,39 @@ export default async function writeScenarioComponents({
             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}`;
+          }
         }
       }
     }
@@ -1202,9 +1492,11 @@ export default async function writeScenarioComponents({
     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
@@ -1316,11 +1608,29 @@ export default async function writeScenarioComponents({
           '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
           // This handles cases where multiple entities are imported from the same index file