@codeyam/codeyam-cli 0.1.0-staging.b8a55ba → 0.1.0-staging.d0ad4ae

package/analyzer-template/packages/analyze/src/lib/files/scenarios/gatherDataForMocks.ts

@@ -3,6 +3,7 @@ import {
   convertDotNotation,
   fillInDirectSchemaGapsAndUnknowns,
   joinParenthesesAndArrays,
+  mergeJsonTypeDefinitions,
   splitOutsideParenthesesAndArrays,
 } from '~codeyam/ai';
 import { DataStructure, JsonTypeDefinition } from '~codeyam/types';
@@ -15,6 +16,65 @@ function cleanFunctionName(name: string): string {
   return name?.split('<')[0];
 }
 
+/**
+ * Escapes special regex characters in a string.
+ */
+function escapeRegExp(str: string): string {
+  return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+}
+
+/**
+ * Counts only DIRECT function calls, excluding method chains on return values.
+ *
+ * When a hook returns an object with methods that are later called, the analysis
+ * tracks those as call signatures like 'useAuth().functionCallReturnValue.signOut()'.
+ * These are NOT separate calls to useAuth - they're method invocations on the return value.
+ *
+ * This function filters out such method chains to get the true count of direct calls.
+ *
+ * @example
+ * // One direct call, one method on return value
+ * countDirectCalls(['useAuth()', 'useAuth().functionCallReturnValue.signOut()']) // => 1
+ *
+ * // Two direct calls
+ * countDirectCalls(['useFetcher()', 'useFetcher()']) // => 2
+ */
+function countDirectCalls(calls: string[]): number {
+  return calls.filter((sig) => !sig.includes('.functionCallReturnValue.'))
+    .length;
+}
+
+/**
+ * Checks if all direct calls (non-method-chain calls) are identical.
+ * This happens when destructuring a hook call:
+ *   const { authLoading, signOut } = useAuth();
+ * produces calls: ['useAuth()', 'useAuth()'] (same signature repeated)
+ *
+ * In this case, we should use the full unqualified schema since there's
+ * really only ONE function call, not multiple different calls.
+ *
+ * @example
+ * // Same call repeated (from destructuring)
+ * allDirectCallsIdentical(['useAuth()', 'useAuth()']) // => true
+ *
+ * // Different calls
+ * allDirectCallsIdentical(['useFetcher<User>()', 'useFetcher<Report>()']) // => false
+ *
+ * // Single call
+ * allDirectCallsIdentical(['useAuth()']) // => true
+ *
+ * // With method chain (method chain is ignored)
+ * allDirectCallsIdentical(['useAuth()', 'useAuth().functionCallReturnValue.signOut()']) // => true
+ */
+function allDirectCallsIdentical(calls: string[]): boolean {
+  const directCalls = calls.filter(
+    (sig) => !sig.includes('.functionCallReturnValue.'),
+  );
+  if (directCalls.length === 0) return true;
+  const firstCall = directCalls[0];
+  return directCalls.every((call) => call === firstCall);
+}
+
 /**
  * Preprocesses the schema to:
  * 1. Filter out all signature paths
@@ -108,18 +168,23 @@ function processCall(
   hookCallIndex: number,
 ): { key: string; value: JsonTypeDefinition } | null {
   const callName = importedExport.calls?.[callIndex] ?? importedExport.name;
+
   // Determine if we should use variable-qualified schema or base schema.
   // Use variable-qualified schema when:
-  // 1. variableNameOccurrence is defined (we're processing multiple calls with matching variable names), OR
-  // 2. calls.length equals callVariableNames.length (single call with single variable name)
+  // 1. variableNameOccurrence is defined (we're processing multiple DISTINCT calls), OR
+  // 2. Single call with single variable name (directCallCount === 1 && varNamesLength === 1)
   // Use base schema when:
-  // - calls.length !== callVariableNames.length (rest spread or method chain case)
+  // - Multiple calls that are all identical (destructuring case: const { a, b } = useHook())
+  // - directCallCount !== varNamesLength (rest spread or method chain case)
   // This is important for rest spread destructuring like: const { x, y, ...data } = useLoaderData()
   // where data.queueState.jobs should come from the base schema, not from a variable-qualified schema.
-  const callsLength = importedExport.calls?.length ?? 0;
+  // Also important: method chains like useAuth().functionCallReturnValue.signOut() should NOT count
+  // as separate calls - only direct hook calls count.
+  const directCallCount = countDirectCalls(importedExport.calls ?? []);
   const varNamesLength = importedExport.callVariableNames?.length ?? 0;
   const shouldUseVariableQualifiedSchema =
-    variableNameOccurrence !== undefined || callsLength === varNamesLength;
+    variableNameOccurrence !== undefined ||
+    (directCallCount === 1 && varNamesLength === 1);
   const variableName = shouldUseVariableQualifiedSchema
     ? importedExport.callVariableNames?.[callIndex]
     : undefined;
@@ -134,15 +199,30 @@ function processCall(
       ? `${variableName}[${variableNameOccurrence}]`
       : variableName;
 
-  const schemaLookupKey = indexedVariableName
-    ? `${importedExport.name}::${indexedVariableName}`
-    : importedExport.name;
+  // Schema lookup order (most specific to least specific):
+  // 1. Call signature (e.g., 'useLoaderData<typeof loader>()')
+  // 2. Variable-qualified key with index (e.g., 'useFetcher::fetcher[1]')
+  // 3. Variable-qualified key without index (e.g., 'useFetcher::fetcher')
+  // 4. Base name (e.g., 'useLoaderData')
 
-  let rawSchema =
-    dependencySchemas?.[importedExport.filePath]?.[schemaLookupKey]
-      ?.returnValueSchema ?? {};
+  let rawSchema: Record<string, string> = {};
+
+  // 1. Try call signature first (most specific - new format)
+  if (callName !== importedExport.name) {
+    rawSchema =
+      dependencySchemas?.[importedExport.filePath]?.[callName]
+        ?.returnValueSchema ?? {};
+  }
+
+  // 2. Try variable-qualified key with index
+  if (Object.keys(rawSchema).length === 0 && indexedVariableName) {
+    const indexedKey = `${importedExport.name}::${indexedVariableName}`;
+    rawSchema =
+      dependencySchemas?.[importedExport.filePath]?.[indexedKey]
+        ?.returnValueSchema ?? {};
+  }
 
-  // If no schema found with indexed key, try without the index as fallback
+  // 3. Try variable-qualified key without index (fallback for indexed lookup)
   if (
     Object.keys(rawSchema).length === 0 &&
     variableNameOccurrence !== undefined &&
@@ -155,15 +235,53 @@ function processCall(
         ?.returnValueSchema ?? {};
   }
 
-  // If no schema found with variable-qualified key, fall back to base name
-  const fallbackSchema =
-    Object.keys(rawSchema).length === 0
-      ? (dependencySchemas?.[importedExport.filePath]?.[importedExport.name]
-          ?.returnValueSchema ?? {})
-      : rawSchema;
+  // 4. Get base name schema and merge with call-specific schema
+  // The base schema often has more complete array element details that the
+  // call-specific schema might be missing. Instead of using as fallback only
+  // when empty, MERGE both schemas so we get the most complete picture.
+  //
+  // HOWEVER: For type-parameterized calls where the base schema contains paths
+  // from MULTIPLE DIFFERENT type parameters, DON'T merge to avoid contamination.
+  // The call-specific schema (rawSchema) is already isolated and complete.
+  const baseSchema =
+    dependencySchemas?.[importedExport.filePath]?.[importedExport.name]
+      ?.returnValueSchema ?? {};
+
+  // Check if this is a type-parameterized call with a call-specific schema
+  const hasTypeParams = callName.includes('<') && callName.includes('>');
+  const hasCallSpecificSchema = Object.keys(rawSchema).length > 0;
+
+  // Check if the base schema has multiple DIFFERENT type parameter variants
+  // If so, merging would bring in wrong fields from other variants
+  let baseHasMultipleTypeParamVariants = false;
+  if (hasTypeParams && hasCallSpecificSchema) {
+    const baseFunctionName = cleanFunctionName(callName.split('(')[0]);
+    const typeParamVariantsInBase = new Set<string>();
+
+    for (const path of Object.keys(baseSchema)) {
+      // Look for paths like "useFetcher<TypeA>(...)" or "useFetcher<TypeB>(...)"
+      const match = path.match(
+        new RegExp(`^${escapeRegExp(baseFunctionName)}<([^>]+)>\\(`),
+      );
+      if (match) {
+        typeParamVariantsInBase.add(match[1]);
+      }
+    }
+    baseHasMultipleTypeParamVariants = typeParamVariantsInBase.size > 1;
+  }
+
+  // For type-parameterized calls where the base schema has MULTIPLE different type
+  // parameter variants, use ONLY the call-specific schema to avoid contamination.
+  // Otherwise, merge with base schema for complete coverage (e.g., array element details).
+  const mergedSchema =
+    hasCallSpecificSchema && hasTypeParams && baseHasMultipleTypeParamVariants
+      ? rawSchema // Use ONLY the call-specific schema when base has mixed type params
+      : Object.keys(rawSchema).length === 0
+        ? baseSchema
+        : { ...baseSchema, ...rawSchema };
 
   // Preprocess to filter signatures and extract functionCallReturnValue
-  let preprocessedSchema = preprocessSchemaForMocks(fallbackSchema);
+  let preprocessedSchema = preprocessSchemaForMocks(mergedSchema);
 
   const functionNames: string[] = [];
 
@@ -297,29 +415,11 @@ function processCall(
     ? `${baseSchemaKey}::${variableName}`
     : baseSchemaKey;
 
-  // Build output key
-  // When entityName is provided, use canonical key format for consistency across scenarios
-  // Otherwise, use the variable-based format for backward compatibility
-  let outputKey: string;
-  if (options?.entityName) {
-    // Canonical key format: EntityName::hookName::index
-    // This ensures consistent keys regardless of variable names chosen during analysis
-    const hookName = cleanFunctionName(importedExport.name);
-    outputKey = `${options.entityName}::${hookName}::${hookCallIndex}`;
-  } else if (variableName) {
-    // Legacy format: "variableName <- functionName"
-    // This reads naturally as "variableName receives from functionName"
-    // When a variable name is reused (variableNameOccurrence > 0), disambiguate with index
-    const disambiguator =
-      variableNameOccurrence !== undefined && variableNameOccurrence > 0
-        ? `[${variableNameOccurrence}]`
-        : '';
-    outputKey = `${variableName}${disambiguator} <- ${importedExport.name}`;
-  } else if (isObjectStyle) {
-    outputKey = importedExport.name;
-  } else {
-    outputKey = callName;
-  }
+  // Build output key - use the ORIGINAL CALL SIGNATURE from the code
+  // This makes it clear to the LLM exactly which call needs which data.
+  // Example: useFetcher<{ data: UserData }>() vs useFetcher<{ data: ReportData }>()
+  // Example: db.select(usersQuery) vs db.select(postsQuery)
+  const outputKey = callName;
 
   // Try to find return value - check with variable qualifier first, then without
   let returnValue = schema[schemaKey];
@@ -334,6 +434,34 @@ function processCall(
     returnValue = schema['functionCallReturnValue'];
   }
 
+  // Handle method-chain style calls (e.g., createClient().from("workouts").select("*") or db.select(query))
+  // After convertDotNotation, the schema is nested: { "createClient()": { "from(...)": { ... } } }
+  // We need to navigate the nested structure using the call signature parts.
+  // For object-style dependencies (isObjectStyle=true), the baseSchemaKey is just the object name (e.g., "db"),
+  // so returnValue might be the parent object, not the actual return value of the method chain.
+  // We should always try method-chain navigation for calls with dots and parens.
+  if (callName.includes('.') && callName.includes('(')) {
+    const callParts = splitOutsideParenthesesAndArrays(callName);
+    let current: unknown = schema;
+    for (const part of callParts) {
+      if (
+        current &&
+        typeof current === 'object' &&
+        !Array.isArray(current) &&
+        part in (current as Record<string, unknown>)
+      ) {
+        current = (current as Record<string, unknown>)[part];
+      } else {
+        current = undefined;
+        break;
+      }
+    }
+    // Use method-chain result if found, otherwise keep previous returnValue
+    if (current !== undefined) {
+      returnValue = current as typeof returnValue;
+    }
+  }
+
   if (!returnValue) {
     return null;
   }
@@ -341,6 +469,27 @@ function processCall(
   return { key: outputKey, value: returnValue as JsonTypeDefinition };
 }
 
+/**
+ * Adds a result to the accumulator, merging schemas if the key already exists.
+ * This handles the case where the same call signature appears multiple times
+ * (e.g., two calls to useFetcher() with the same generic type).
+ */
+function addResultToAccumulator(
+  acc: { [key: string]: JsonTypeDefinition },
+  key: string,
+  value: JsonTypeDefinition,
+): void {
+  if (acc[key]) {
+    // Key already exists - merge the schemas
+    acc[key] = mergeJsonTypeDefinitions(
+      acc[key] as Record<string, unknown>,
+      value as Record<string, unknown>,
+    ) as JsonTypeDefinition;
+  } else {
+    acc[key] = value;
+  }
+}
+
 export default function gatherDataForMocks(
   importedExports: {
     filePath?: string;
@@ -363,14 +512,121 @@ export default function gatherDataForMocks(
       // Get the clean hook name for indexing
       const hookName = cleanFunctionName(importedExport.name);
 
-      // Check if we have multiple calls with variable names
-      const hasMultipleCalls =
-        importedExport.calls &&
-        importedExport.calls.length > 1 &&
+      // Check if we have multiple DISTINCT direct calls with variable names.
+      // Important: Only count direct calls to the function, not method chains on its return value.
+      // For example, ['useAuth()', 'useAuth().functionCallReturnValue.signOut()'] has only 1 direct call.
+      //
+      // ALSO: If all direct calls are IDENTICAL (same signature repeated), this is from
+      // destructuring like `const { a, b } = useHook()` and we should use the FULL unqualified
+      // schema, not iterate through partial variable-qualified schemas.
+      const directCallCount = countDirectCalls(importedExport.calls ?? []);
+      const hasMultipleDistinctCalls =
+        directCallCount > 1 &&
         importedExport.callVariableNames &&
-        importedExport.callVariableNames.length === importedExport.calls.length;
+        importedExport.callVariableNames.length === directCallCount &&
+        !allDirectCallsIdentical(importedExport.calls ?? []);
+
+      // Check for multiple distinct method-chain calls WITHOUT callVariableNames.
+      // This pattern is common with tRPC, Supabase, etc. where a single import
+      // (e.g., `trpc`) is used with multiple different method chains (e.g.,
+      // `trpc.useUtils()`, `trpc.quote.getQuotesByOrg.useQuery()`).
+      // These should each get their own entry in dataForMocks.
+      //
+      // IMPORTANT: This is DIFFERENT from the rest spread pattern where a function
+      // like useLoaderData is called directly, and other calls are method chains
+      // on its return value (e.g., `useLoaderData<typeof loader>()` with
+      // `useLoaderData<typeof loader>().uncommittedEntitiesList.filter(...)`).
+      // In that case, all the additional calls start with the base function call.
+      const directCalls = (importedExport.calls ?? []).filter(
+        (sig) => !sig.includes('.functionCallReturnValue.'),
+      );
+      const uniqueDirectCalls = [...new Set(directCalls)];
+
+      // Check if any call is a direct invocation of the import (e.g., `useLoaderData<typeof loader>()`)
+      // as opposed to all being property accesses (e.g., `trpc.useUtils()`, `trpc.quote.xxx()`).
+      // If there's a direct function call, the other calls are likely method chains on its result.
+      const baseName = importedExport.name;
+      const hasDirectFunctionCall = directCalls.some((sig) => {
+        // Pattern: baseName() or baseName<...>()
+        // This matches "useLoaderData()", "useLoaderData<typeof loader>()", etc.
+        // but NOT "trpc.useUtils()" which has a dot before the parens
+        const baseCallPattern = new RegExp(`^${baseName}(<[^>]*>)?\\(`);
+        return baseCallPattern.test(sig);
+      });
+
+      const hasMultipleDistinctMethodChains =
+        !hasMultipleDistinctCalls &&
+        uniqueDirectCalls.length > 1 &&
+        !allDirectCallsIdentical(directCalls) &&
+        !hasDirectFunctionCall; // Exclude rest spread patterns
+
+      if (hasMultipleDistinctMethodChains) {
+        // For multiple method-chain calls on the same import, create a NESTED structure
+        // under the base import name. This allows constructMockCode to generate a
+        // combined mock object with all method chains.
+        //
+        // Example: For calls ['trpc.useUtils()', 'trpc.quote.getQuotesByOrg.useQuery()'],
+        // create: { "trpc": { "useUtils()": {...}, "quote": { "getQuotesByOrg": { "useQuery()": {...} } } } }
+        const nestedResult: Record<string, unknown> = {};
+
+        for (let i = 0; i < uniqueDirectCalls.length; i++) {
+          // Find the index of this call in the original calls array
+          const callIndex = (importedExport.calls ?? []).indexOf(
+            uniqueDirectCalls[i],
+          );
+          if (callIndex === -1) continue;
 
-      if (hasMultipleCalls) {
+          const hookCallIndex = hookCallIndices[hookName] ?? 0;
+          hookCallIndices[hookName] = hookCallIndex + 1;
+
+          const processResult = processCall(
+            importedExport,
+            dependencySchemas,
+            callIndex,
+            undefined,
+            options,
+            hookCallIndex,
+          );
+
+          if (processResult) {
+            // The key is like "trpc.useUtils()" or "trpc.quote.getQuotesByOrg.useQuery()"
+            // We need to nest the value at the path AFTER the base name
+            const callSignature = processResult.key;
+
+            // Split the call signature into path parts, handling parens
+            // e.g., "trpc.useUtils()" -> ["trpc", "useUtils()"]
+            // e.g., "trpc.quote.getQuotesByOrg.useQuery()" -> ["trpc", "quote", "getQuotesByOrg", "useQuery()"]
+            const pathParts = splitOutsideParenthesesAndArrays(callSignature);
+
+            // Skip the base name (first part) and nest the value at the remaining path
+            if (pathParts.length > 1 && pathParts[0] === baseName) {
+              const nestedPath = pathParts.slice(1);
+
+              // Set the value at the nested path, merging with existing structure
+              let current = nestedResult;
+              for (let j = 0; j < nestedPath.length - 1; j++) {
+                const part = nestedPath[j];
+                if (!(part in current)) {
+                  current[part] = {};
+                }
+                current = current[part] as Record<string, unknown>;
+              }
+              // Set the leaf value
+              const leafKey = nestedPath[nestedPath.length - 1];
+              current[leafKey] = processResult.value;
+            }
+          }
+        }
+
+        // Add the nested result under the base import name
+        if (Object.keys(nestedResult).length > 0) {
+          addResultToAccumulator(
+            acc,
+            baseName,
+            nestedResult as JsonTypeDefinition,
+          );
+        }
+      } else if (hasMultipleDistinctCalls) {
         // Track how many times each variable name has been seen
         // to disambiguate keys when the same variable is reused
         const variableNameCounts: Record<string, number> = {};
@@ -396,15 +652,16 @@ export default function gatherDataForMocks(
             hookCallIndex,
           );
           if (processResult) {
-            acc[processResult.key] = processResult.value;
+            addResultToAccumulator(acc, processResult.key, processResult.value);
           }
         }
       } else {
-        // Original behavior for single calls or calls without variable names
+        // Single calls, or multiple identical calls (from destructuring)
         // Get the current hook call index and increment for next call
         const hookCallIndex = hookCallIndices[hookName] ?? 0;
         hookCallIndices[hookName] = hookCallIndex + 1;
 
+        // First, try to get result from unqualified schema
         const processResult = processCall(
           importedExport,
           dependencySchemas,
@@ -413,8 +670,98 @@ export default function gatherDataForMocks(
           options,
           hookCallIndex,
         );
-        if (processResult) {
-          acc[processResult.key] = processResult.value;
+
+        // Check if we got a meaningful result (not empty object or null)
+        // Arrays are always meaningful (even empty ones represent valid data structure)
+        const hasUnqualifiedResult =
+          processResult &&
+          processResult.value !== null &&
+          (typeof processResult.value !== 'object' ||
+            Array.isArray(processResult.value) ||
+            Object.keys(processResult.value as object).length > 0);
+
+        // Check if we have multiple callVariableNames that are DIFFERENT
+        // (e.g., from parent + child components using the same hook differently)
+        // In this case, we should also merge in variable-qualified schemas
+        // to capture all destructured properties from both parent and child.
+        const uniqueCallVarNames = new Set(
+          importedExport.callVariableNames ?? [],
+        );
+        const hasDifferentVarNames =
+          importedExport.callVariableNames &&
+          importedExport.callVariableNames.length > 1 &&
+          uniqueCallVarNames.size > 1 &&
+          // Only merge variable-qualified schemas when calls are all identical
+          // (e.g., useFlags() repeated). Don't do this for method chains
+          // or rest spread cases where calls.length !== callVariableNames.length
+          allDirectCallsIdentical(importedExport.calls ?? []);
+
+        if (hasUnqualifiedResult) {
+          let finalValue = processResult.value;
+
+          // If we have different variable names (e.g., from parent and child components),
+          // also merge in variable-qualified schemas to ensure we capture all properties.
+          // This handles the case where parent uses useFlags().enableNewAdminDashboard
+          // and child uses useFlags().flagShowReferenceIdOnQuoteTable.
+          if (hasDifferentVarNames) {
+            let mergedValue = processResult.value as JsonTypeDefinition;
+            for (let i = 0; i < importedExport.callVariableNames!.length; i++) {
+              const varQualifiedResult = processCall(
+                importedExport,
+                dependencySchemas,
+                i,
+                0, // Use 0 to trigger variable-qualified lookup
+                options,
+                hookCallIndex,
+              );
+              if (varQualifiedResult && varQualifiedResult.value) {
+                mergedValue = mergeJsonTypeDefinitions(
+                  mergedValue as Record<string, unknown>,
+                  varQualifiedResult.value as Record<string, unknown>,
+                ) as JsonTypeDefinition;
+              }
+            }
+            finalValue = mergedValue;
+          }
+
+          addResultToAccumulator(
+            acc,
+            processResult.key,
+            finalValue as JsonTypeDefinition,
+          );
+        } else if (
+          directCallCount > 1 &&
+          allDirectCallsIdentical(importedExport.calls ?? []) &&
+          importedExport.callVariableNames &&
+          importedExport.callVariableNames.length > 0
+        ) {
+          // No unqualified schema exists, but we have multiple identical calls with variable names
+          // Merge all the variable-qualified schemas together
+          const callSignature =
+            importedExport.calls?.[0] ?? `${importedExport.name}()`;
+          let mergedValue: JsonTypeDefinition = {};
+
+          for (let i = 0; i < importedExport.callVariableNames.length; i++) {
+            const varQualifiedResult = processCall(
+              importedExport,
+              dependencySchemas,
+              i,
+              0, // Use 0 to trigger variable-qualified lookup
+              options,
+              hookCallIndex,
+            );
+            if (varQualifiedResult && varQualifiedResult.value) {
+              // Deep merge the value into the accumulated result
+              mergedValue = mergeJsonTypeDefinitions(
+                mergedValue,
+                varQualifiedResult.value,
+              );
+            }
+          }
+
+          if (Object.keys(mergedValue as object).length > 0) {
+            addResultToAccumulator(acc, callSignature, mergedValue);
+          }
         }
       }