codegen.ts

import type { CodegenPlugin } from '@graphql-codegen/plugin-helpers';

import { plugin as typescriptPlugin } from '@graphql-codegen/typescript';
import { format } from 'prettier';
import ts from 'typescript';

import { Create, gray } from '@syncify/ansi';

interface CustomPluginConfig {
  export: string[];
  depths?: number;
  outputDir?: string;
  skipTypename?: boolean;
  commentDescriptions?: boolean;
}

const customPlugin: CodegenPlugin<CustomPluginConfig> = {
  plugin: async (schema, documents, config, info) => {
    const tsConfig = {
      skipTypename: config.skipTypename ?? true,
      commentDescriptions: config.commentDescriptions ?? true
    };
    const typescriptOutput = await typescriptPlugin(schema, documents, tsConfig, info);
    return processPluginOutput(typescriptOutput.content, config);
  }
};

async function processPluginOutput (content: string, config: CustomPluginConfig) {
  const log = Create().Break().Top('Generating Types').Newline()
  .toLog({ clear: true });
  const { transformedContent, typeLiterals, includedTypes } = await transformTypes(content, config, log);

  if (typeLiterals.length > 0 && config.outputDir) {
    logOptionalTypeLiterals(typeLiterals, config.outputDir, log);
  }

  logEnd(includedTypes.size, log);
  return { content: transformedContent };
}

async function transformTypes (content: string, config: CustomPluginConfig, log: any) {
  const sourceFile = ts.createSourceFile('temp.ts', content, ts.ScriptTarget.Latest, true);
  const { typeMap, enumNames, typeLiterals } = collectTypes(sourceFile, log);
  const includedTypes = collectIncludedTypes(typeMap, config, log);
  const transformedNodes = transformNodes(typeMap, includedTypes, enumNames, sourceFile, log);
  let fileContent = printNodes(transformedNodes, sourceFile, log);

  fileContent = replaceScalars(fileContent, log);
  fileContent = cleanUpContent(fileContent, log);

  const transformedContent = await format('/* eslint-disable no-use-before-define */\n\n' + fileContent, {
    parser: 'typescript',
    tabWidth: 2,
    printWidth: 120,
    singleQuote: true,
    useTabs: false
  });

  return { transformedContent, typeLiterals, includedTypes };
}

function collectTypes (sourceFile: ts.SourceFile, log: any) {
  const typeMap = new Map<string, ts.TypeAliasDeclaration | ts.InterfaceDeclaration>();
  const enumNames = new Set<string>();
  const typeLiterals: string[] = [];

  sourceFile.forEachChild((node) => {
    if (ts.isTypeAliasDeclaration(node) || ts.isInterfaceDeclaration(node)) {
      const name = node.name?.text;
      if (name) {
        typeMap.set(name, node);
        typeLiterals.push(name);
        log.Line(`Found type: ${name} (${ts.isInterfaceDeclaration(node) ? 'interface' : 'type alias'})`);
      }
    } else if (ts.isEnumDeclaration(node) && node.name) {
      enumNames.add(node.name.text);
      log.Line(`Found enum: ${node.name.text} - will be replaced with 'string'`);
    }
  });

  return { typeMap, enumNames, typeLiterals };
}

function collectIncludedTypes (typeMap: Map<string, ts.TypeAliasDeclaration | ts.InterfaceDeclaration>, config: CustomPluginConfig, log: any) {
  const includedTypes = new Map<string, number>();
  config.export.forEach((type) => includedTypes.set(type, 0));

  const collectReferencedTypes = (node: ts.Node, depth: number) => {
    ts.forEachChild(node, (child) => {
      if (ts.isTypeReferenceNode(child) && child.typeName && ts.isIdentifier(child.typeName)) {
        const typeName = child.typeName.text;
        if (typeMap.has(typeName) && !isBuiltInType(typeName)) {
          const currentDepth = includedTypes.get(typeName);
          if (currentDepth === undefined || currentDepth > depth + 1) {
            includedTypes.set(typeName, depth + 1);
            log.Line(`Including ${typeName} at depth ${depth + 1}`);
            if (depth + 1 < (config.depths || Infinity) || config.export.includes(typeName)) {
              collectReferencedTypes(typeMap.get(typeName)!, depth + 1);
            }
          }
        }
      } else if (ts.isArrayTypeNode(child)) {
        collectReferencedTypes(child.elementType, depth);
      } else if (ts.isUnionTypeNode(child) || ts.isIntersectionTypeNode(child)) {
        child.types.forEach((type) => collectReferencedTypes(type, depth));
      } else if (ts.isParenthesizedTypeNode(child)) {
        collectReferencedTypes(child.type, depth);
      }
      collectReferencedTypes(child, depth);
    });
  };

  config.export.forEach((type) => {
    if (typeMap.has(type)) collectReferencedTypes(typeMap.get(type)!, 0);
    else log.Warn(`Type ${type} not found in source`);
  });

  return includedTypes;
}

function transformNodes (
  typeMap: Map<string, ts.TypeAliasDeclaration | ts.InterfaceDeclaration>,
  includedTypes: Map<string, number>,
  enumNames: Set<string>,
  sourceFile: ts.SourceFile,
  log: any
) {
  const transformedNodes: ts.Node[] = [];

  const visitType = (node: ts.TypeNode): ts.TypeNode => {
    log.Line(`Visiting type node: ${node.getText(sourceFile)}`);
    if (ts.isTypeReferenceNode(node) && ts.isIdentifier(node.typeName)) {
      const refType = node.typeName.text;
      if (refType === 'Maybe' && node.typeArguments?.length) {
        log.Line(`Unwrapping Maybe<${node.typeArguments[0].getText(sourceFile)}>`);
        return visitType(node.typeArguments[0]);
      } else if (enumNames.has(refType)) {
        log.Line(`Replacing enum ${refType} with 'string'`);
        return ts.factory.createKeywordTypeNode(ts.SyntaxKind.StringKeyword);
      }
    }

    if (ts.isArrayTypeNode(node)) {
      const elementType = visitType(node.elementType);
      if (elementType !== node.elementType) return ts.factory.createArrayTypeNode(elementType);
    } else if (ts.isUnionTypeNode(node)) {
      const types = node.types.map(visitType);
      const allAny = types.every((t) => t.kind === ts.SyntaxKind.AnyKeyword);
      const hasNonAny = types.some((t) => t.kind !== ts.SyntaxKind.AnyKeyword);
      if (allAny) return ts.factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword);
      if (hasNonAny) {
        const cleanedTypes = types.filter((t) => t.kind !== ts.SyntaxKind.AnyKeyword);
        if (cleanedTypes.length === types.length) return node;
        if (cleanedTypes.length === 0) return ts.factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword);
        // Explicitly cast to ts.TypeNode[] to satisfy createUnionTypeNode
        return ts.factory.createUnionTypeNode([ ...cleanedTypes, ts.factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword) ] as ts.TypeNode[]);
      }
    }
    return node;
  };

  includedTypes.forEach((_, type) => {
    const node = typeMap.get(type);
    if (!node) {
      log.Warn(`Type ${type} not found in typeMap`);
      return;
    }
    transformedNodes.push(transformNode(node, type, visitType, sourceFile, log));
  });

  augmentQueryRoot(typeMap.get('QueryRoot'), transformedNodes, visitType, sourceFile, log);
  return finalizeNodes(transformedNodes, includedTypes, sourceFile, log);
}

function transformNode (node: ts.TypeAliasDeclaration | ts.InterfaceDeclaration, type: string, visitType: (node: ts.TypeNode) => ts.TypeNode, sourceFile: ts.SourceFile, log: any) {
  let transformedNode: ts.TypeAliasDeclaration | ts.InterfaceDeclaration;

  if (type.endsWith('Payload')) {
    log.Line(`Transforming Payload type: ${type}`);
    const prefix = type.replace(/Payload$/, '');
    const lowerPrefix = prefix.charAt(0).toLowerCase() + prefix.slice(1);
    const newTypeName = `Mutation${prefix}`;

    const jsDoc = getJsDoc(node, sourceFile);
    const originalType = ts.isTypeAliasDeclaration(node) ? node.type : ts.factory.createTypeLiteralNode(node.members);
    const wrappedType = ts.factory.createTypeLiteralNode([
      ts.factory.createPropertySignature(undefined, lowerPrefix, undefined, visitType(originalType))
    ]);

    transformedNode = ts.factory.createTypeAliasDeclaration(
      [ ts.factory.createModifier(ts.SyntaxKind.ExportKeyword) ],
      ts.factory.createIdentifier(newTypeName),
      undefined,
      wrappedType
    );

    if (jsDoc) {
      transformedNode = ts.addSyntheticLeadingComment(transformedNode, ts.SyntaxKind.MultiLineCommentTrivia, `*\n * ${jsDoc}\n `, true);
    }
  } else {
    transformedNode = ts.isTypeAliasDeclaration(node)
      ? ts.factory.createTypeAliasDeclaration(
        node.modifiers ?? [ ts.factory.createModifier(ts.SyntaxKind.ExportKeyword) ],
        node.name,
        node.typeParameters,
        visitType(node.type)
      )
      : ts.factory.createInterfaceDeclaration(
        node.modifiers ?? [ ts.factory.createModifier(ts.SyntaxKind.ExportKeyword) ],
        node.name,
        node.typeParameters,
        node.heritageClauses,
        node.members.map((member) =>
          ts.isPropertySignature(member) && member.type
            ? ts.factory.createPropertySignature(member.modifiers, member.name, member.questionToken, visitType(member.type))
            : member)
      );

    const jsDoc = getJsDoc(node, sourceFile);
    if (jsDoc) {
      transformedNode = ts.addSyntheticLeadingComment(transformedNode, ts.SyntaxKind.MultiLineCommentTrivia, `*\n * ${jsDoc}\n `, true);
    }
  }

  return transformedNode;
}

function augmentQueryRoot (queryRootNode: ts.TypeAliasDeclaration | ts.InterfaceDeclaration | undefined, transformedNodes: ts.Node[], visitType: (node: ts.TypeNode) => ts.TypeNode, sourceFile: ts.SourceFile, log: any) {
  if (!queryRootNode) {
    log.Warn('QueryRoot not found in typeMap');
    return;
  }

  const processProperty = (prop: ts.PropertySignature) => {
    if (prop.name && ts.isIdentifier(prop.name) && prop.type) {
      const propName = prop.name.getText(sourceFile);
      const typeName = `Query${propName.charAt(0).toUpperCase() + propName.slice(1)}`;
      const jsDoc = getJsDoc(prop, sourceFile);
      const innerType = visitType(prop.type);
      const objectType = ts.factory.createTypeLiteralNode([
        ts.factory.createPropertySignature(
          undefined,
          propName,
          undefined,
          innerType.kind === ts.SyntaxKind.TypeReference &&
          ts.isIdentifier((innerType as ts.TypeReferenceNode).typeName) &&
          (innerType as ts.TypeReferenceNode).typeName.getText(sourceFile) === 'Maybe' &&
          (innerType as ts.TypeReferenceNode).typeArguments?.length
            ? (innerType as ts.TypeReferenceNode).typeArguments[0]
            : innerType
        )
      ]);

      const finalType =
        innerType.kind === ts.SyntaxKind.TypeReference &&
        ts.isIdentifier((innerType as ts.TypeReferenceNode).typeName) &&
        (innerType as ts.TypeReferenceNode).typeName.getText(sourceFile) === 'Maybe'
          ? ts.factory.createTypeReferenceNode('Maybe', [ objectType ])
          : objectType;

      let newTypeAlias = ts.factory.createTypeAliasDeclaration(
        [ ts.factory.createModifier(ts.SyntaxKind.ExportKeyword) ],
        ts.factory.createIdentifier(typeName),
        undefined,
        finalType
      );

      if (jsDoc) {
        newTypeAlias = ts.addSyntheticLeadingComment(newTypeAlias, ts.SyntaxKind.MultiLineCommentTrivia, `*\n * ${jsDoc}\n `, true);
      }

      transformedNodes.push(newTypeAlias);
      log.Line(`Generated ${typeName} for QueryRoot property ${propName}`);
    } else {
      log.Line(`Skipping property in QueryRoot: ${prop.name ? prop.name.getText(sourceFile) : 'unnamed'} - missing required fields`, gray);
    }
  };

  if (ts.isInterfaceDeclaration(queryRootNode)) {
    log.Line('Processing QueryRoot (interface) for augmentation');
    queryRootNode.members.filter(ts.isPropertySignature).forEach(processProperty);
  } else if (ts.isTypeAliasDeclaration(queryRootNode) && ts.isTypeLiteralNode(queryRootNode.type)) {
    log.Line('Processing QueryRoot (type alias) for augmentation');
    queryRootNode.type.members.filter(ts.isPropertySignature).forEach(processProperty);
  } else {
    log.Line('QueryRoot type is not a type literal or interface, skipping augmentation');
  }
}

function finalizeNodes (transformedNodes: ts.Node[], includedTypes: Map<string, number>, sourceFile: ts.SourceFile, log: any) {
  const definedTypes = new Set(includedTypes.keys());
  const final: ts.Node[] = [];

  transformedNodes.forEach((node) => {
    const finalNode = ts.transform(node, [
      (context) => (root) => {
        const visit = (node: ts.Node): ts.Node => {
          if (ts.isTypeReferenceNode(node) && ts.isIdentifier(node.typeName)) {
            const refType = node.typeName.text;
            if (!definedTypes.has(refType) && !isBuiltInType(refType)) {
              log.Line(`Post-processing: Replacing undefined ${refType} with 'any'`);
              return ts.factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword);
            }
          } else if (ts.isArrayTypeNode(node) && ts.isTypeReferenceNode(node.elementType) && ts.isIdentifier(node.elementType.typeName)) {
            const elemTypeName = node.elementType.typeName.text;
            if (!definedTypes.has(elemTypeName) && !isBuiltInType(elemTypeName)) {
              log.Line(`Post-processing: Replacing Array<${elemTypeName}> with Array<any>`);
              return ts.factory.createArrayTypeNode(ts.factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword));
            }
          } else if (ts.isUnionTypeNode(node)) {
            const types = node.types.map(visit);
            const allAny = types.every((t) => t.kind === ts.SyntaxKind.AnyKeyword);
            const hasNonAny = types.some((t) => t.kind !== ts.SyntaxKind.AnyKeyword);
            if (allAny) return ts.factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword);
            if (hasNonAny) {
              const cleanedTypes = types.filter((t) => t.kind !== ts.SyntaxKind.AnyKeyword);
              if (cleanedTypes.length === types.length) return node;
              if (cleanedTypes.length === 0) return ts.factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword);
              return ts.factory.createUnionTypeNode([ ...cleanedTypes, ts.factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword) ] as ts.TypeNode[]);
            }
          }
          return ts.visitEachChild(node, visit, context);
        };
        return ts.visitNode(root, visit) as typeof root;
      }
    ]).transformed[0];
    final.push(finalNode);
  });

  return final.filter((node) => {
    if (ts.isTypeAliasDeclaration(node)) {
      const name = node.name.text;
      if (name === 'InputMaybe') {
        log.Line('Removing InputMaybe type definition');
        return false;
      } else if (name === 'Scalars') {
        log.Line('Removing Scalars type definition');
        return false;
      }
    }
    return true;
  });
}

function printNodes (nodes: ts.Node[], sourceFile: ts.SourceFile, log: any) {
  const printer = ts.createPrinter({ newLine: ts.NewLineKind.LineFeed });
  return nodes.map((node) => printer.printNode(ts.EmitHint.Unspecified, node, sourceFile)).join('\n\n');
}

function replaceScalars (content: string, log: any) {
  const scalarMapping: Record<string, string> = {
    ID: 'string',
    String: 'string',
    URL: 'string',
    Date: 'string',
    DateTime: 'string',
    Color: 'string',
    UnsignedInt64: 'string',
    UtcOffset: 'string',
    Decimal: 'string',
    Money: 'string',
    FormattedString: 'string',
    HTML: 'string',
    Int: 'number',
    Float: 'number',
    BigInt: 'number',
    Boolean: 'boolean',
    ARN: 'any',
    StorefrontID: 'any',
    JSON: 'Record<string, any>'
  };

  return content.replace(/Scalars\['([^']+)'\]\['(input|output)'\]/g, (match, scalarName) => {
    const replacement = scalarMapping[scalarName] || 'any';
    log.Line(`Replacing ${match} with ${replacement}`);
    return replacement;
  });
}

function cleanUpContent (content: string, log: any) {
  return content
  .replace(/\* \*/g, '*')
  .replace(/(\/\*\*.*?\*\/)\s*\1/g, '$1')
  .replace(/InputMaybe<([^]+?)>(?=;)/g, '$1')
  .replace(/export type Maybe<T> = null \| any;/, '')
  .replace(/(?<=(?:<|: ))Maybe<([^]+?)>(?=(?:;|>))/g, '$1');
}

function getJsDoc (node: ts.Node, sourceFile: ts.SourceFile) {
  const leadingComments = ts.getLeadingCommentRanges(sourceFile.text, node.getFullStart());
  if (leadingComments) {
    const jsDocComment = leadingComments
    .filter((comment) => comment.kind === ts.SyntaxKind.MultiLineCommentTrivia)
    .map((comment) => sourceFile.text.slice(comment.pos + 2, comment.end - 2).trim())[0];
    return jsDocComment?.replace(/^\/\*\*|\*\/$/g, '').trim();
  }
  return undefined;
}

function logOptionalTypeLiterals (typeLiterals: string[], outputDir: string, log: any) {
  const uniqueLiterals = Array.from(new Set(typeLiterals)).map((name) => `'${name}'`).join(' | ');
  const dtsContent = [
    'import type { LiteralUnion } from \'type-fest\';',
    '',
    `export type Models = LiteralUnion<${uniqueLiterals}, string>;`
  ].join('\n');
  console.warn(`Generated codegen.d.ts content:\n${dtsContent}\nUse a CLI hook to write this separately.`);
}

function logEnd (typeCount: number, log: any) {
  log.Line(`Extracted ${typeCount} types (including dependencies)`).Newline().End('Generated Types').Break()
  .toLog({ clear: true });
}

function isBuiltInType (typeName: string) {
  return [ 'string', 'number', 'boolean', 'null', 'undefined', 'any', 'never', 'unknown', 'Array' ].includes(typeName);
}

export default customPlugin;