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

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

@@ -1,8 +1,10 @@
-// import { writeFile } from 'fs/promises';
 import {
   convertDotNotation,
   fillInDirectSchemaGapsAndUnknowns,
+  buildSchemaIndexes,
+  type SchemaIndexes,
   joinParenthesesAndArrays,
+  mergeJsonTypeDefinitions,
   splitOutsideParenthesesAndArrays,
 } from '~codeyam/ai';
 import { DataStructure, JsonTypeDefinition } from '~codeyam/types';
@@ -15,6 +17,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
@@ -68,131 +129,808 @@ function preprocessSchemaForMocks(
   return result;
 }
 
+/**
+ * Options for gatherDataForMocks
+ */
+export interface GatherDataForMocksOptions {
+  /**
+   * When provided, generates canonical keys in format `EntityName::hookName::index`.
+   * This ensures consistent keys across scenarios regardless of variable names
+   * chosen during analysis.
+   */
+  entityName?: string;
+}
+
+/**
+ * Processes a single call (with optional variable name) and returns the mock data entry.
+ * Returns null if no valid schema is found.
+ *
+ * @param importedExport - The imported export to process
+ * @param dependencySchemas - The dependency schemas
+ * @param callIndex - The index of the call to process
+ * @param variableNameOccurrence - When a variable name is reused, this is the occurrence index
+ *                                  (e.g., first 'data' = 0, second 'data' = 1)
+ *                                  Used to disambiguate keys when the same variable is reused.
+ * @param options - Additional options including entityName for canonical key generation
+ * @param hookCallIndex - The global index of this hook call for canonical key generation
+ */
+function processCall(
+  importedExport: {
+    filePath?: string;
+    name: string;
+    isMocked?: boolean;
+    calls?: string[];
+    callVariableNames?: string[];
+  },
+  dependencySchemas: DataStructure['dependencySchemas'],
+  callIndex: number,
+  variableNameOccurrence: number | undefined,
+  options: GatherDataForMocksOptions | undefined,
+  hookCallIndex: number,
+  prebuiltIndexes?: SchemaIndexes,
+): { 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 DISTINCT calls), OR
+  // 2. Single call with single variable name (directCallCount === 1 && varNamesLength === 1)
+  // Use base schema when:
+  // - 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.
+  // 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 ||
+    (directCallCount === 1 && varNamesLength === 1);
+  const variableName = shouldUseVariableQualifiedSchema
+    ? importedExport.callVariableNames?.[callIndex]
+    : undefined;
+
+  // Build the schema lookup key - use variable-qualified key if available
+  // When the same variable name is reused (variableNameOccurrence > 0), try indexed lookup first
+  // e.g., "useFetcher::fetcher[1]" for the second use of "fetcher"
+  const indexedVariableName =
+    variableName &&
+    variableNameOccurrence !== undefined &&
+    variableNameOccurrence > 0
+      ? `${variableName}[${variableNameOccurrence}]`
+      : variableName;
+
+  // 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: 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 ?? {};
+  }
+
+  // 3. Try variable-qualified key without index (fallback for indexed lookup)
+  if (
+    Object.keys(rawSchema).length === 0 &&
+    variableNameOccurrence !== undefined &&
+    variableNameOccurrence > 0 &&
+    variableName
+  ) {
+    const nonIndexedKey = `${importedExport.name}::${variableName}`;
+    rawSchema =
+      dependencySchemas?.[importedExport.filePath]?.[nonIndexedKey]
+        ?.returnValueSchema ?? {};
+  }
+
+  // 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 we have a call-specific schema, we need to
+  // be careful about merging with the base schema. The base schema may contain
+  // merged data from ALL type parameter variants, which would pollute the result.
+  let mergedSchema: Record<string, string>;
+  if (hasCallSpecificSchema && hasTypeParams) {
+    if (baseHasMultipleTypeParamVariants) {
+      // Base schema has paths with different type parameters - filter to only
+      // include paths matching the current type parameter.
+      const typeParamMatch = callName.match(/<([^>]+)>/);
+      const currentTypeParam = typeParamMatch ? typeParamMatch[1] : null;
+      const baseFunctionName = cleanFunctionName(callName.split('(')[0]);
+
+      if (currentTypeParam) {
+        // Filter base schema to only include paths that:
+        // 1. Match current type param (e.g., useLoaderData<LoaderData>())
+        // 2. OR don't have any type param (generic paths like useLoaderData())
+        const filteredBaseSchema: Record<string, string> = {};
+        for (const [path, type] of Object.entries(baseSchema)) {
+          // Check if path has a type parameter for this function
+          const pathTypeParamMatch = path.match(
+            new RegExp(`^${escapeRegExp(baseFunctionName)}<([^>]+)>\\(`),
+          );
+          if (pathTypeParamMatch) {
+            // Path has a type param - only include if it matches current
+            if (pathTypeParamMatch[1] === currentTypeParam) {
+              filteredBaseSchema[path] = type;
+            }
+            // Otherwise skip (it's for a different type param variant)
+          } else {
+            // Path doesn't have a type param for this function - include it
+            // (could be generic paths or paths for other functions)
+            filteredBaseSchema[path] = type;
+          }
+        }
+        mergedSchema = { ...filteredBaseSchema, ...rawSchema };
+      } else {
+        // Couldn't extract type param - fall back to rawSchema only
+        mergedSchema = rawSchema;
+      }
+    } else {
+      // Base schema has zero or one type-parameterized variant.
+      // When there's only one variant, the base schema paths are safe to merge
+      // (no contamination from other variants). This is important for cases like:
+      // - Array element details (gitStatus[].path) that are in base but not call-specific
+      // - Enriched paths that should be preserved
+      // Filter to only include paths matching current type param or generic paths.
+      const typeParamMatch = callName.match(/<([^>]+)>/);
+      const currentTypeParam = typeParamMatch ? typeParamMatch[1] : null;
+      const baseFunctionName = cleanFunctionName(callName.split('(')[0]);
+
+      if (currentTypeParam) {
+        const filteredBaseSchema: Record<string, string> = {};
+        for (const [path, type] of Object.entries(baseSchema)) {
+          const pathTypeParamMatch = path.match(
+            new RegExp(`^${escapeRegExp(baseFunctionName)}<([^>]+)>\\(`),
+          );
+          if (pathTypeParamMatch) {
+            // Path has a type param - only include if it matches current
+            if (pathTypeParamMatch[1] === currentTypeParam) {
+              filteredBaseSchema[path] = type;
+            }
+          } else {
+            // Path doesn't have a type param - include generic paths
+            filteredBaseSchema[path] = type;
+          }
+        }
+        mergedSchema = { ...filteredBaseSchema, ...rawSchema };
+      } else {
+        // No type param to extract - merge normally
+        mergedSchema = { ...baseSchema, ...rawSchema };
+      }
+    }
+  } else if (Object.keys(rawSchema).length === 0) {
+    mergedSchema = baseSchema;
+  } else {
+    mergedSchema = { ...baseSchema, ...rawSchema };
+  }
+
+  // Preprocess to filter signatures and extract functionCallReturnValue
+  let preprocessedSchema = preprocessSchemaForMocks(mergedSchema);
+
+  const isFunctionNameWithoutFunction = (
+    path: string,
+    functionName: string,
+  ) => {
+    // Check if path starts with functionName but isn't a function call
+    // Must also handle generics: useLoaderData<typeof loader>() should NOT
+    // be filtered when functionName is 'useLoaderData'
+    return (
+      path.startsWith(functionName) &&
+      !path.startsWith(functionName + '(') &&
+      !path.startsWith(functionName + '.(') &&
+      !path.startsWith(functionName + '<') // Handle generics
+    );
+  };
+
+  // For variable-qualified schemas, the paths also include the variable qualifier
+  // e.g., "useFetcher()::entityDiffFetcher.functionCallReturnValue.data"
+  // We need to handle both qualified and non-qualified paths
+
+  // Merge non-generic variants into generic variants
+  // e.g., if callName is "useFetcher<T>()", also include paths from "useFetcher()"
+  if (callName.includes('<') && callName.includes('(')) {
+    // Extract the base function name without generics
+    const baseFunctionName = cleanFunctionName(callName.split('(')[0]);
+    const nonGenericCall = variableName
+      ? `${baseFunctionName}()::${variableName}`
+      : baseFunctionName + '()';
+    const genericCallPrefix = variableName
+      ? `${callName.endsWith(')') ? callName : callName + '()'}::${variableName}`
+      : callName.endsWith(')')
+        ? callName
+        : callName + '()';
+
+    // If we have a generic call, merge paths from non-generic variant
+    const mergedSchema: Record<string, string> = { ...preprocessedSchema };
+    for (const path of Object.keys(preprocessedSchema)) {
+      if (path.startsWith(nonGenericCall)) {
+        // Create corresponding generic path
+        const genericPath =
+          genericCallPrefix + path.slice(nonGenericCall.length);
+        // Only add if the generic path doesn't already exist
+        if (!mergedSchema[genericPath]) {
+          mergedSchema[genericPath] = preprocessedSchema[path];
+        }
+      }
+    }
+    preprocessedSchema = mergedSchema;
+  }
+
+  // Merge generic call (e.g., "foo()") data into specific argument calls (e.g., "foo(arg1, arg2)")
+  // This handles the case where the generic schema has richer nested data than the specific call
+  if (callName.includes('(') && !callName.includes('<')) {
+    // Extract the base function name and create generic call
+    const baseFunctionName = callName.split('(')[0];
+
+    // For path matching, use non-qualified call names since schema paths don't include
+    // variable qualifiers. Variable qualifiers are only for schema lookup, not path matching.
+    const genericCallForMatching = baseFunctionName + '()';
+    const specificCallForMatching = callName;
+
+    // Only merge if the callName has arguments (not just empty parens)
+    if (specificCallForMatching !== genericCallForMatching) {
+      const mergedSchema: Record<string, string> = {
+        ...preprocessedSchema,
+      };
+      for (const path of Object.keys(preprocessedSchema)) {
+        if (path.startsWith(genericCallForMatching)) {
+          // Create corresponding specific path
+          const specificPath =
+            specificCallForMatching + path.slice(genericCallForMatching.length);
+          // Add if the specific path doesn't exist, OR if it exists but is "unknown"
+          // The generic call's schema often has more precise type info than the specific call
+          if (
+            !mergedSchema[specificPath] ||
+            mergedSchema[specificPath] === 'unknown'
+          ) {
+            mergedSchema[specificPath] = preprocessedSchema[path];
+          }
+        }
+      }
+      preprocessedSchema = mergedSchema;
+    }
+  }
+
+  // Two-pass function name filtering to avoid O(k²) retroactive deletion.
+  // Pass 1: Discover all unique function names from paths ending with ')'
+  const functionNameSet = new Set<string>();
+  const schemaKeys = Object.keys(preprocessedSchema);
+  for (const path of schemaKeys) {
+    if (path.endsWith(')')) {
+      const pathParts = splitOutsideParenthesesAndArrays(path);
+      const functionName = joinParenthesesAndArrays([
+        ...pathParts.slice(0, -1),
+        pathParts[pathParts.length - 1].split('(')[0],
+      ]);
+      functionNameSet.add(functionName);
+    }
+  }
+
+  // Index function names by first segment for O(1) bucket lookup instead of
+  // checking all function names for every path. Most paths only need to check
+  // 1-2 function names instead of all F.
+  const fnByFirstSegment = new Map<string, string[]>();
+  for (const fn of functionNameSet) {
+    const dotIdx = fn.indexOf('.');
+    const firstSeg = dotIdx >= 0 ? fn.slice(0, dotIdx) : fn;
+    let bucket = fnByFirstSegment.get(firstSeg);
+    if (!bucket) {
+      bucket = [];
+      fnByFirstSegment.set(firstSeg, bucket);
+    }
+    bucket.push(fn);
+  }
+
+  // Pass 2: Filter paths that match any discovered function name prefix.
+  // Use indexed lookup: extract path's first segment, check only matching function names.
+  const relevantMergedDependencySchema: Record<string, string> = {};
+  for (const path of schemaKeys) {
+    // Extract first segment of path (before first '.', '(', or '<')
+    const dotIdx = path.indexOf('.');
+    const parenIdx = path.indexOf('(');
+    const angleIdx = path.indexOf('<');
+    let minIdx = path.length;
+    if (dotIdx >= 0 && dotIdx < minIdx) minIdx = dotIdx;
+    if (parenIdx >= 0 && parenIdx < minIdx) minIdx = parenIdx;
+    if (angleIdx >= 0 && angleIdx < minIdx) minIdx = angleIdx;
+    const firstSeg = path.slice(0, minIdx);
+
+    const candidates = fnByFirstSegment.get(firstSeg);
+    if (
+      !candidates ||
+      !candidates.some((fn) => isFunctionNameWithoutFunction(path, fn))
+    ) {
+      relevantMergedDependencySchema[path] = preprocessedSchema[path];
+    }
+  }
+
+  const filledSchema = fillInDirectSchemaGapsAndUnknowns({
+    scopeName: importedExport.name,
+    schema: relevantMergedDependencySchema,
+    prebuiltIndexes,
+  });
+
+  const schema = convertDotNotation(filledSchema);
+
+  // Determine if this is an object-style or hook-style dependency:
+  // - Hook-style: calls = ['useAuth()'] - call starts with name + '('
+  // - Object-style: calls = ['supabase.from(...)'] - call starts with name + '.'
+  const isObjectStyle =
+    callName.startsWith(importedExport.name + '.') ||
+    (!callName.includes('(') && callName === importedExport.name);
+
+  // Build schema key with variable qualifier if present
+  const baseSchemaKey = isObjectStyle ? importedExport.name : callName;
+  const schemaKey = variableName
+    ? `${baseSchemaKey}::${variableName}`
+    : baseSchemaKey;
+
+  // 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];
+  if (!returnValue && variableName) {
+    returnValue = schema[baseSchemaKey];
+  }
+  if (!returnValue) {
+    returnValue = schema['returnValue'];
+  }
+  // Fallback for function call return values (e.g., useFetcher returns via functionCallReturnValue)
+  if (!returnValue) {
+    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;
+  }
+
+  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;
     name: string;
     isMocked?: boolean;
     calls?: string[];
+    callVariableNames?: string[];
   }[],
   dependencySchemas: DataStructure['dependencySchemas'],
+  options?: GatherDataForMocksOptions,
 ) {
-  return importedExports.reduce(
+  // Build schema indexes ONCE from all dependency schemas combined.
+  // This avoids rebuilding indexes (~1-2 seconds for 18k keys) for each processCall.
+  //
+  // IMPORTANT: We must preprocess schemas BEFORE building indexes!
+  // The schemas contain .functionCallReturnValue paths that get stripped during processing.
+  // If we build indexes from original paths, lookups with preprocessed paths will fail.
+  // Example: 'hook().functionCallReturnValue.x.length' → index maps 'hook().functionCallReturnValue.x' → ['length']
+  //          But preprocessed path 'hook().x.length' looks for 'hook().x' which doesn't exist in the index.
+  const allSchemaKeys: Record<string, string> = {};
+  if (dependencySchemas) {
+    for (const filePath in dependencySchemas) {
+      for (const entityName in dependencySchemas[filePath]) {
+        const schema =
+          dependencySchemas[filePath][entityName]?.returnValueSchema;
+        if (schema) {
+          // Preprocess each schema to remove .functionCallReturnValue before collecting
+          const preprocessed = preprocessSchemaForMocks(schema);
+          Object.assign(allSchemaKeys, preprocessed);
+        }
+      }
+    }
+  }
+  const prebuiltIndexes =
+    Object.keys(allSchemaKeys).length > 0
+      ? buildSchemaIndexes(allSchemaKeys)
+      : undefined;
+
+  // Track the global index for each hook name across all imports
+  // This enables canonical keys like EntityName::useLoaderData::0, EntityName::useFetcher::0, etc.
+  const hookCallIndices: Record<string, number> = {};
+
+  const result = importedExports.reduce(
     (acc, importedExport) => {
       if (!importedExport.isMocked) return acc;
 
-      const functionNames: string[] = [];
-
-      const isFunctionNameWithoutFunction = (
-        path: string,
-        functionName: string,
-      ) => {
-        // Check if path starts with functionName but isn't a function call
-        // Must also handle generics: useLoaderData<typeof loader>() should NOT
-        // be filtered when functionName is 'useLoaderData'
-        return (
-          path.startsWith(functionName) &&
-          !path.startsWith(functionName + '(') &&
-          !path.startsWith(functionName + '.(') &&
-          !path.startsWith(functionName + '<') // Handle generics
-        );
-      };
+      // Get the clean hook name for indexing
+      const hookName = cleanFunctionName(importedExport.name);
+
+      // 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 === 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;
 
-      const rawSchema =
-        dependencySchemas?.[importedExport.filePath]?.[importedExport.name]
-          ?.returnValueSchema ?? {};
-
-      // Preprocess to filter signatures and extract functionCallReturnValue
-      let preprocessedSchema = preprocessSchemaForMocks(rawSchema);
-
-      // Merge non-generic variants into generic variants
-      // e.g., if callName is "useFetcher<T>()", also include paths from "useFetcher()"
-      const callName = importedExport.calls?.[0] ?? importedExport.name;
-      if (callName.includes('<') && callName.includes('(')) {
-        // Extract the base function name without generics
-        const baseFunctionName = cleanFunctionName(callName.split('(')[0]);
-        const nonGenericCall = baseFunctionName + '()';
-        const genericCallPrefix = callName.endsWith(')')
-          ? callName
-          : callName + '()';
-
-        // If we have a generic call, merge paths from non-generic variant
-        const mergedSchema: Record<string, string> = { ...preprocessedSchema };
-        for (const path of Object.keys(preprocessedSchema)) {
-          if (path.startsWith(nonGenericCall)) {
-            // Create corresponding generic path
-            const genericPath =
-              genericCallPrefix + path.slice(nonGenericCall.length);
-            // Only add if the generic path doesn't already exist
-            if (!mergedSchema[genericPath]) {
-              mergedSchema[genericPath] = preprocessedSchema[path];
+          const hookCallIndex = hookCallIndices[hookName] ?? 0;
+          hookCallIndices[hookName] = hookCallIndex + 1;
+
+          const processResult = processCall(
+            importedExport,
+            dependencySchemas,
+            callIndex,
+            undefined,
+            options,
+            hookCallIndex,
+            prebuiltIndexes,
+          );
+
+          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;
             }
           }
         }
-        preprocessedSchema = mergedSchema;
-      }
 
-      const relevantMergedDependencySchema = Object.keys(
-        preprocessedSchema,
-      ).reduce(
-        (acc, path) => {
-          if (path.endsWith(')')) {
-            const pathParts = splitOutsideParenthesesAndArrays(path);
-            const functionName = joinParenthesesAndArrays([
-              ...pathParts.slice(0, -1),
-              pathParts[pathParts.length - 1].split('(')[0],
-            ]);
-            functionNames.push(functionName);
-            for (const existingPath in acc) {
-              if (isFunctionNameWithoutFunction(existingPath, functionName)) {
-                delete acc[existingPath];
-              }
-            }
+        // 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> = {};
+
+        // Process each call separately with its variable name
+        for (let i = 0; i < importedExport.calls!.length; i++) {
+          const varName = importedExport.callVariableNames![i];
+
+          // Calculate the occurrence index for this variable name
+          const occurrence = variableNameCounts[varName] ?? 0;
+          variableNameCounts[varName] = occurrence + 1;
+
+          // Get the current hook call index and increment for next call
+          const hookCallIndex = hookCallIndices[hookName] ?? 0;
+          hookCallIndices[hookName] = hookCallIndex + 1;
+
+          const processResult = processCall(
+            importedExport,
+            dependencySchemas,
+            i,
+            occurrence,
+            options,
+            hookCallIndex,
+            prebuiltIndexes,
+          );
+          if (processResult) {
+            addResultToAccumulator(acc, processResult.key, processResult.value);
           }
+        }
+      } else if (
+        // Handle multiple UNIQUE type-parameterized direct calls when variable count doesn't match.
+        // This happens when a hook is called with different type parameters (useFetcher<TypeA>(), useFetcher<TypeB>()),
+        // but the callVariableNames count doesn't match directCallCount (common with rest spread patterns).
+        // Each unique call signature should get its own entry in dataForMocks.
+        uniqueDirectCalls.length > 1 &&
+        !allDirectCallsIdentical(directCalls) &&
+        hasDirectFunctionCall
+      ) {
+        // Process each UNIQUE direct call signature
+        for (let i = 0; i < uniqueDirectCalls.length; i++) {
+          const callSignature = uniqueDirectCalls[i];
 
-          if (
-            functionNames.some((fn) => isFunctionNameWithoutFunction(path, fn))
-          ) {
-            return acc;
+          // Find the index of this call in the original calls array
+          const callIndex = (importedExport.calls ?? []).indexOf(callSignature);
+          if (callIndex === -1) continue;
+
+          const hookCallIndex = hookCallIndices[hookName] ?? 0;
+          hookCallIndices[hookName] = hookCallIndex + 1;
+
+          const processResult = processCall(
+            importedExport,
+            dependencySchemas,
+            callIndex,
+            undefined,
+            options,
+            hookCallIndex,
+            prebuiltIndexes,
+          );
+
+          if (processResult) {
+            addResultToAccumulator(acc, processResult.key, processResult.value);
           }
+        }
+      } else {
+        // 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;
 
-          acc[path] = preprocessedSchema[path];
-          return acc;
-        },
-        {} as Record<string, string>,
-      );
+        // First, try to get result from unqualified schema
+        const processResult = processCall(
+          importedExport,
+          dependencySchemas,
+          0,
+          undefined,
+          options,
+          hookCallIndex,
+          prebuiltIndexes,
+        );
 
-      const schema = convertDotNotation(
-        fillInDirectSchemaGapsAndUnknowns({
-          scopeName: importedExport.name,
-          schema: relevantMergedDependencySchema,
-        }),
-      );
+        // 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);
 
-      // Determine if this is an object-style or hook-style dependency:
-      // - Hook-style: calls = ['useAuth()'] - call starts with name + '('
-      // - Object-style: calls = ['supabase.from(...)'] - call starts with name + '.'
-      //
-      // After convertDotNotation:
-      // - Hook-style schema has the callName as top-level key: schema['useAuth()']
-      // - Object-style schema has the import name as top-level key: schema['supabase']
-      // Note: callName is already defined above for generic/non-generic merging
-      const isObjectStyle =
-        callName.startsWith(importedExport.name + '.') ||
-        (!callName.includes('(') && callName === importedExport.name);
-
-      const schemaKey = isObjectStyle ? importedExport.name : callName;
-      const outputKey = isObjectStyle ? importedExport.name : callName;
-      const returnValue = schema[schemaKey] ?? schema['returnValue'];
-
-      if (!returnValue) {
-        return acc;
-      }
+        // 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,
+                prebuiltIndexes,
+              );
+              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,
+              prebuiltIndexes,
+            );
+            if (varQualifiedResult && varQualifiedResult.value) {
+              // Deep merge the value into the accumulated result
+              mergedValue = mergeJsonTypeDefinitions(
+                mergedValue,
+                varQualifiedResult.value,
+              );
+            }
+          }
 
-      acc[outputKey] = returnValue as JsonTypeDefinition;
+          if (Object.keys(mergedValue as object).length > 0) {
+            addResultToAccumulator(acc, callSignature, mergedValue);
+          }
+        }
+      }
 
       return acc;
     },
     {} as { [importAlias: string]: JsonTypeDefinition },
   );
+
+  return result;
 }