From 3b12b442143067caf2a4a5bd6acd02a88b575288 Mon Sep 17 00:00:00 2001
From: Chris <chrisbreuer93@gmail.com>
Date: Mon, 4 Nov 2024 17:29:29 +0100
Subject: [PATCH] chore: wip

---
 fixtures/output/function.d.ts |  2 +-
 src/extract.ts                | 50 +++++++++++++++++++++++++----------
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/fixtures/output/function.d.ts b/fixtures/output/function.d.ts
index a236a44..12927d4 100644
--- a/fixtures/output/function.d.ts
+++ b/fixtures/output/function.d.ts
@@ -13,4 +13,4 @@ export declare function processData(data: unknown): unknown;
 export declare function complexAsyncGenerator(): any;
 export declare function isUser(value: unknown): value is User;
 export declare function extractFunctionSignature(declaration: string): FunctionSignature;
-export declare function createApi<T extends Record<string, (...args: any[]) => any>>(endpoints: T): { [K in keyof T]: ReturnType<T[K]> extends Promise<infer R> ? R : ReturnType<T[K]> } { return {} as any };
+export declare function createApi<T extends Record<string, (...args: any[]) => any>>(endpoints: T): { [K in keyof T]: ReturnType<T[K]> extends Promise<infer R> ? R : ReturnType<T[K]> };
diff --git a/src/extract.ts b/src/extract.ts
index 94c60aa..6db3648 100644
--- a/src/extract.ts
+++ b/src/extract.ts
@@ -422,55 +422,77 @@ function extractParams(rest: string): { params: string, rest: string } {
 function extractReturnType(rest: string, declaration: string): { returnType: string } {
   let returnType = 'void'
   if (rest.startsWith(':')) {
-    debugLog(undefined, 'signature-return-start', `Processing return type from: ${rest}`)
+    debugLog(undefined, 'return-start', `Starting return type extraction with: ${rest}`)
     rest = rest.slice(1).trim()
 
-    if (!rest) {
-      // If rest is empty, perhaps the return type is on the next line(s)
-      const indexOfColon = declaration.indexOf(':')
-      rest = declaration.slice(indexOfColon + 1).trim()
-    }
-
     let depth = 0
     let buffer = ''
     let i = 0
     let inString = false
     let stringChar = ''
+    let foundEnd = false
+
+    debugLog(undefined, 'return-extraction', 'Starting character-by-character extraction')
 
-    while (i < rest.length) {
+    while (i < rest.length && !foundEnd) {
       const char = rest[i]
       const prevChar = i > 0 ? rest[i - 1] : ''
+      const nextChar = i < rest.length - 1 ? rest[i + 1] : ''
 
-      // Handle string literals
+      debugLog(undefined, 'return-char',
+        `Pos ${i}: Char "${char}", Depth ${depth}, InString ${inString}, Buffer length ${buffer.length}`)
+
+      // Handle string boundaries
       if ((char === '"' || char === '\'' || char === '`') && prevChar !== '\\') {
         if (!inString) {
           inString = true
           stringChar = char
+          debugLog(undefined, 'return-string', `Entering string with ${stringChar}`)
         }
         else if (char === stringChar) {
           inString = false
+          debugLog(undefined, 'return-string', 'Exiting string')
         }
       }
 
       // Track depth when not in string
       if (!inString) {
-        if (char === '{' || char === '<' || char === '(')
+        if (char === '{' || char === '<' || char === '(') {
           depth++
-        else if (char === '}' || char === '>' || char === ')')
+          debugLog(undefined, 'return-depth', `Opening bracket, increasing depth to ${depth}`)
+        }
+        else if (char === '}' || char === '>' || char === ')') {
           depth--
+          debugLog(undefined, 'return-depth', `Closing bracket, decreasing depth to ${depth}`)
+
+          // If we hit depth 0 with a closing brace, this might be the end of our type
+          if (depth === 0 && char === '}') {
+            buffer += char
+            // Look ahead to see if this is followed by a function body
+            const nextNonWhitespace = rest.slice(i + 1).trim()[0]
+            if (nextNonWhitespace === '{') {
+              debugLog(undefined, 'return-end', `Found end of return type at pos ${i}, next char is function body`)
+              foundEnd = true
+              break
+            }
+          }
+        }
 
-        // Stop at function body start or when we hit a semicolon outside any depth
-        if ((depth === 0 && char === '{') || (depth === 0 && char === ';')) {
+        // Stop at semicolons at depth 0
+        if (depth === 0 && char === ';') {
+          debugLog(undefined, 'return-end', 'Found semicolon at depth 0')
+          foundEnd = true
           break
         }
       }
 
       buffer += char
+      debugLog(undefined, 'return-buffer', `Updated buffer: ${buffer}`)
       i++
     }
 
     returnType = buffer.trim()
-    debugLog(undefined, 'signature-return', `Extracted return type: ${returnType}`)
+    debugLog(undefined, 'return-final', `Final extracted return type: ${returnType}`)
   }
   return { returnType }
 }