alias_generator.go

package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"log/slog"
	"os"
	"sort"
	"strings"
	"unicode"

	"golang.org/x/tools/go/packages"
)

func generateAliases(inputFilePath, outputFilePath string) error {
	outputPackageNameSplit := strings.Split(outputFilePath, "/")
	outputPackageName := outputPackageNameSplit[len(outputPackageNameSplit)-2]

	// Parse the input file
	fset := token.NewFileSet()

	file, err := parser.ParseFile(fset, inputFilePath, nil, parser.ParseComments)
	if err != nil {
		return fmt.Errorf("parse input file: %w", err)
	}

	outputStringBuilder := &strings.Builder{}

	fmt.Fprintf(outputStringBuilder, "// Code generated by github.com/raito-io/sdk-go/generator, DO NOT EDIT.\n\n")

	pkg, err := packages.Load(&packages.Config{Mode: packages.NeedName})
	if err != nil {
		return fmt.Errorf("load package information: %w", err)
	}

	filePathSplit := strings.Split(inputFilePath, "/")
	filePath := strings.Join(filePathSplit[:len(filePathSplit)-1], "/")
	pkgAlias := filePathSplit[len(filePathSplit)-2]

	pkgPath := pkg[0].PkgPath + "/" + filePath

	constAssignments := make(map[string][]string)

	// Write package declaration
	fmt.Fprintf(outputStringBuilder, "package %s\n\n", outputPackageName)

	// Write import statements
	fmt.Fprintf(outputStringBuilder, "import %s \"%s\"\n", pkgAlias, pkgPath)

	// Iterate through the file's declarations
	for _, decl := range file.Decls {
		if d, ok := decl.(*ast.GenDecl); ok {
			if d.Tok == token.TYPE {
				// Iterate through type specs
				for _, spec := range d.Specs {
					if ts, ok := spec.(*ast.TypeSpec); ok {
						// Check if it is a struct
						if unicode.IsUpper(rune(ts.Name.Name[0])) {
							// Write alias declaration
							fmt.Fprintf(outputStringBuilder, "type %s = %s.%s\n", ts.Name, pkgAlias, ts.Name)
						}
					}
				}
			} else if d.Tok == token.CONST {
				var assignments []string

				// Iterator through const specs
				for _, spec := range d.Specs {
					if vs, ok := spec.(*ast.ValueSpec); ok {
						if it, ok := vs.Type.(*ast.Ident); ok {
							for i := range vs.Names {
								if unicode.IsUpper(rune(vs.Names[i].Name[0])) {
									constAssignments[it.Name] = append(constAssignments[it.Name], vs.Names[i].Name)

									assignments = append(assignments, fmt.Sprintf("%s %s = %s.%s", vs.Names[i].Name, it.Name, pkgAlias, vs.Names[i].Name))
								}
							}
						}
					}
				}

				if len(assignments) > 0 {
					fmt.Fprintf(outputStringBuilder, "const (%s\n)\n\n", strings.Join(assignments, "\n\t"))
				}
			}
		}
	}

	if len(constAssignments) > 0 {
		fmt.Fprintf(outputStringBuilder, "\n")

		keys := make([]string, 0, len(constAssignments))
		for k := range constAssignments {
			keys = append(keys, k)
		}

		sort.Strings(keys)

		for _, constType := range keys {
			assignments := constAssignments[constType]
			fmt.Fprintf(outputStringBuilder, "var All%s = []%s{%s}\n", constType, constType, strings.Join(assignments, ", "))
		}
	}

	// Format the output file
	formattedOutput, err := format.Source([]byte(outputStringBuilder.String()))
	if err != nil {
		return fmt.Errorf("format generated file: %w", err)
	}

	// Create the output file
	outputFile, err := os.Create(outputFilePath)
	if err != nil {
		return fmt.Errorf("create generated file: %w", err)
	}
	defer outputFile.Close()

	// Write the formatted output to the file
	_, err = outputFile.Write(formattedOutput)
	if err != nil {
		return fmt.Errorf("write generated file: %w", err)
	}

	return nil
}

func main() {
	// Parse command line flags
	inputFile := flag.String("input", "", "Input Go file path")
	outputFile := flag.String("output", "", "Output Go file path")
	flag.Parse()

	// Validate input parameters
	if *inputFile == "" || *outputFile == "" {
		slog.Error("Usage: go run main.go -input <input_file_path> -package <output_package_name> -output <output_file_path>")
		os.Exit(1)
	}

	// Generate aliases
	err := generateAliases(*inputFile, *outputFile)
	if err != nil {
		slog.Error("Error:", err)
		os.Exit(1)
	}

	slog.Info("Code generation completed successfully.")
}