From dc1734f4bb457e383c579cca7b40a1682cc3bf9d Mon Sep 17 00:00:00 2001
From: Jon Ambas <jon@jonambas.com>
Date: Sun, 7 Apr 2024 12:53:46 -0400
Subject: [PATCH] feat(codegen): add ct return value hints (#16)

---
 src/__tests__/codegen.test.ts |  9 ++++++++-
 src/codegen.ts                | 17 ++++++++---------
 src/utils.ts                  |  9 +++++++++
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/src/__tests__/codegen.test.ts b/src/__tests__/codegen.test.ts
index 58368db..b899ad4 100644
--- a/src/__tests__/codegen.test.ts
+++ b/src/__tests__/codegen.test.ts
@@ -37,7 +37,14 @@ describe('codegen', () => {
               "file": "ct.mjs",
             },
             {
-              "code": "export const ct: (alias: "foo.100" | "foo.200" | "bar.100" | "bar.200") => string;",
+              "code": "const pluginCtMap = {
+        'foo.100': '#fff',
+        'foo.200': {"base":"#000","lg":"#111"},
+        'bar.100': 'red',
+        'bar.200': 'blue',
+      } as const;
+          
+      export const ct: <T extends keyof typeof pluginCtMap>(alias: T) => typeof pluginCtMap[T];",
               "file": "ct.d.ts",
             },
           ],
diff --git a/src/codegen.ts b/src/codegen.ts
index 31c51f3..3326581 100644
--- a/src/codegen.ts
+++ b/src/codegen.ts
@@ -6,7 +6,7 @@ import type {
 } from '@pandacss/types';
 import { makePaths, mapTemplate } from './map';
 import type { PluginContext } from './types';
-import { serializeValue } from './utils';
+import { serializeMapTypes, serializeValue } from './utils';
 
 export const codegen = (
   args: CodegenPrepareHookArgs,
@@ -20,20 +20,19 @@ export const codegen = (
 
   if (!cssFn || !indexFile || !indexDtsFile) return args.artifacts;
 
-  const code = `${mapTemplate(map)}
+  const ctFile: ArtifactContent = {
+    file: 'ct.mjs',
+    code: `${mapTemplate(map)}
   export const ct = (path) => {
     if (!pluginCtMap.has(path)) return 'panda-plugin-ct_alias-not-found';
     return pluginCtMap.get(path);
-  };`;
-
-  const paths = makePaths(tokens)
-    .map((key) => `"${key}"`)
-    .join(' | ');
+  };`,
+  };
 
-  const ctFile: ArtifactContent = { file: 'ct.mjs', code };
   const ctDtsFile: ArtifactContent = {
     file: 'ct.d.ts',
-    code: `export const ct: (alias: ${paths}) => string;`,
+    code: `${serializeMapTypes(map)}
+    \nexport const ct: <T extends keyof typeof pluginCtMap>(alias: T) => typeof pluginCtMap[T];`,
   };
 
   cssFn.files.push(ctFile, ctDtsFile);
diff --git a/src/utils.ts b/src/utils.ts
index 1648a93..d482416 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -13,3 +13,12 @@ export const serializeMap = (map: Map<any, any>) => {
 export const serializeValue = (value: any) => {
   return isObject(value) ? JSON.stringify(value) : `'${value}'`;
 };
+
+export const serializeMapTypes = (map: Map<any, any>) => {
+  let code = 'const pluginCtMap = {';
+  for (const [key, value] of map.entries()) {
+    code += `\n  '${key}': ${serializeValue(value)},`;
+  }
+  code += '\n} as const;';
+  return code;
+};