main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"strconv"
)

type operationInfo struct {
	FieldName     string
	OperationName string
}

type serviceInfo struct {
	VarName     string
	ServiceName string
	Operations  []operationInfo
}

func main() {
	var inputPath string
	var outputPath string
	flag.StringVar(&inputPath, "input", "", "input generated proto-model service file")
	flag.StringVar(&outputPath, "output", "", "output registry file")
	flag.Parse()

	if inputPath == "" || outputPath == "" {
		fmt.Fprintln(os.Stderr, "both -input and -output are required")
		os.Exit(2)
	}

	services, err := parseServices(inputPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "parse services: %v\n", err)
		os.Exit(1)
	}
	if len(services) == 0 {
		fmt.Fprintln(os.Stderr, "no nexus services found")
		os.Exit(1)
	}

	source, err := renderRegistry(services)
	if err != nil {
		fmt.Fprintf(os.Stderr, "render registry: %v\n", err)
		os.Exit(1)
	}
	if err := os.WriteFile(outputPath, source, 0o644); err != nil {
		fmt.Fprintf(os.Stderr, "write registry: %v\n", err)
		os.Exit(1)
	}
}

func parseServices(path string) ([]serviceInfo, error) {
	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, path, nil, parser.SkipObjectResolution)
	if err != nil {
		return nil, err
	}

	var services []serviceInfo
	for _, decl := range file.Decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok || genDecl.Tok != token.VAR {
			continue
		}
		for _, spec := range genDecl.Specs {
			valueSpec, ok := spec.(*ast.ValueSpec)
			if !ok || len(valueSpec.Names) != 1 || len(valueSpec.Values) != 1 {
				continue
			}
			compositeLit, ok := valueSpec.Values[0].(*ast.CompositeLit)
			if !ok {
				continue
			}
			structType, ok := valueSpec.Type.(*ast.StructType)
			if !ok {
				structType, ok = compositeLit.Type.(*ast.StructType)
				if !ok {
					continue
				}
			}

			serviceName, ok := compositeServiceName(compositeLit)
			if !ok {
				continue
			}

			var operations []operationInfo
			for _, field := range structType.Fields.List {
				for _, name := range field.Names {
					if name.Name == "ServiceName" {
						continue
					}
					operationName, ok := compositeOperationName(compositeLit, name.Name)
					if !ok {
						return nil, fmt.Errorf("service %s field %s missing operation literal", valueSpec.Names[0].Name, name.Name)
					}
					operations = append(operations, operationInfo{
						FieldName:     name.Name,
						OperationName: operationName,
					})
				}
			}
			if len(operations) == 0 {
				continue
			}
			services = append(services, serviceInfo{
				VarName:     valueSpec.Names[0].Name,
				ServiceName: serviceName,
				Operations:  operations,
			})
		}
	}
	return services, nil
}

func compositeServiceName(lit *ast.CompositeLit) (string, bool) {
	for _, elt := range lit.Elts {
		kv, ok := elt.(*ast.KeyValueExpr)
		if !ok {
			continue
		}
		key, ok := kv.Key.(*ast.Ident)
		if !ok || key.Name != "ServiceName" {
			continue
		}
		return stringLiteral(kv.Value)
	}
	return "", false
}

func compositeOperationName(lit *ast.CompositeLit, fieldName string) (string, bool) {
	for _, elt := range lit.Elts {
		kv, ok := elt.(*ast.KeyValueExpr)
		if !ok {
			continue
		}
		key, ok := kv.Key.(*ast.Ident)
		if !ok || key.Name != fieldName {
			continue
		}
		call, ok := kv.Value.(*ast.CallExpr)
		if !ok || len(call.Args) == 0 {
			return "", false
		}
		return stringLiteral(call.Args[len(call.Args)-1])
	}
	return "", false
}

func stringLiteral(expr ast.Expr) (string, bool) {
	lit, ok := expr.(*ast.BasicLit)
	if !ok || lit.Kind != token.STRING {
		return "", false
	}
	value, err := strconv.Unquote(lit.Value)
	if err != nil {
		return "", false
	}
	return value, true
}

func renderRegistry(services []serviceInfo) ([]byte, error) {
	var buf bytes.Buffer
	buf.WriteString("// Code generated by systemnexusgen. DO NOT EDIT.\n\n")
	buf.WriteString("package workflowservicenexus\n\n")
	buf.WriteString("import (\n")
	buf.WriteString("\t\"reflect\"\n\n")
	buf.WriteString("\t\"google.golang.org/protobuf/proto\"\n")
	buf.WriteString(")\n\n")
	buf.WriteString("type TemporalNexusOperationKey struct {\n")
	buf.WriteString("\tServiceName   string\n")
	buf.WriteString("\tOperationName string\n")
	buf.WriteString("}\n\n")
	buf.WriteString("type TemporalNexusOperationReference interface {\n")
	buf.WriteString("\tName() string\n")
	buf.WriteString("\tInputType() reflect.Type\n")
	buf.WriteString("\tOutputType() reflect.Type\n")
	buf.WriteString("}\n\n")
	buf.WriteString("var TemporalNexusOperationRegistry = map[TemporalNexusOperationKey]TemporalNexusOperationReference{\n")
	for _, service := range services {
		for _, operation := range service.Operations {
			fmt.Fprintf(&buf, "\t{ServiceName: %q, OperationName: %q}: %s.%s,\n", service.ServiceName, operation.OperationName, service.VarName, operation.FieldName)
		}
	}
	buf.WriteString("}\n\n")
	buf.WriteString("func GetTemporalNexusOperation(serviceName, operationName string) TemporalNexusOperationReference {\n")
	buf.WriteString("\treturn TemporalNexusOperationRegistry[TemporalNexusOperationKey{ServiceName: serviceName, OperationName: operationName}]\n")
	buf.WriteString("}\n\n")
	buf.WriteString("func IsTemporalNexusOperation(serviceName, operationName string) bool {\n")
	buf.WriteString("\treturn GetTemporalNexusOperation(serviceName, operationName) != nil\n")
	buf.WriteString("}\n\n")
	buf.WriteString("func NewTemporalNexusOperationInputIfSystemOperation(serviceName, operationName string) proto.Message {\n")
	buf.WriteString("\top := GetTemporalNexusOperation(serviceName, operationName)\n")
	buf.WriteString("\tif op == nil {\n")
	buf.WriteString("\t\treturn nil\n")
	buf.WriteString("\t}\n")
	buf.WriteString("\tinputType := op.InputType()\n")
	buf.WriteString("\tif inputType == nil {\n")
	buf.WriteString("\t\treturn nil\n")
	buf.WriteString("\t}\n")
	buf.WriteString("\tif inputType.Kind() == reflect.Ptr {\n")
	buf.WriteString("\t\tmsg, _ := reflect.New(inputType.Elem()).Interface().(proto.Message)\n")
	buf.WriteString("\t\treturn msg\n")
	buf.WriteString("\t}\n")
	buf.WriteString("\tmsg, _ := reflect.New(inputType).Interface().(proto.Message)\n")
	buf.WriteString("\treturn msg\n")
	buf.WriteString("}\n")

	return format.Source(buf.Bytes())
}