main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"os"
	"strings"

	temporalv1 "github.com/thomas-maurice/protoc-gen-go-tmprl/gen/temporal/v1"
	"github.com/thomas-maurice/protoc-gen-go-tmprl/internal/model"
	"github.com/thomas-maurice/protoc-gen-go-tmprl/internal/renderer"
	"github.com/thomas-maurice/protoc-gen-go-tmprl/internal/version"
	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/pluginpb"
)

var (
	genWorkflowPrefix bool
	genDocs           bool
	// Default activity start to close timeout in seconds
	defaultActivityScheduleToClose int
)

func main() {
	var flags flag.FlagSet
	// This is to generate automatically prefixes for the jobs.
	flags.BoolVar(&genWorkflowPrefix, "gen-workflow-prefix", false, "Generates a prefix for the jobs like foo.v1.Foo.Method/<workflowID>")
	flags.IntVar(&defaultActivityScheduleToClose, "default-activity-schedule-to-close", 3600*24, "Default start to close activity timeout if none is specified anywhere, in seconds")
	flags.BoolVar(&genDocs, "gen-docs", false, "Generates documentation for the temporal workflows")
	opts := &protogen.Options{
		ParamFunc: flags.Set,
	}

	opts.Run(func(gen *protogen.Plugin) error {

		if defaultActivityScheduleToClose <= 0 {
			return fmt.Errorf("the default schedule to close activity timeout cannot be 0 nor negative")
		}

		gen.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)

		config := &model.Config{
			GenWorkflowPrefix:              genWorkflowPrefix,
			DefaultActivityScheduleToClose: defaultActivityScheduleToClose,
		}

		for _, f := range gen.Files {
			if !f.Generate {
				continue
			}

			if err := generateFile(gen, f, config); err != nil {
				return err
			}

			if genDocs {
				if err := generateReadme(gen, f, config); err != nil {
					return err
				}
			}
		}
		return nil
	})
}

func generateFile(plugin *protogen.Plugin, file *protogen.File, config *model.Config) error {
	filename := file.GeneratedFilenamePrefix + "_tmprl.pb.go"

	needsGenerate := false
	for _, s := range file.Services {
		if so, ok := proto.GetExtension(s.Desc.Options(), temporalv1.E_Service).(*temporalv1.ServiceOptions); ok && so != nil {
			needsGenerate = true
		}
	}

	if !needsGenerate {
		return nil
	}

	// A proto package that resolves to the Go package name `temporal` collides
	// with the plugin's own temporal/v1 imports and produces broken code; warn
	// loudly rather than letting the user hit a cryptic buf resolution error.
	if pkg := file.Proto.GetPackage(); pkg == "temporal" || strings.HasPrefix(pkg, "temporal.") {
		fmt.Fprintf(os.Stderr, "protoc-gen-go-tmprl: warning: proto package %q collides with the plugin's own temporal/v1 imports and will generate broken code; rename the package\n", pkg)
	}

	gen := plugin.NewGeneratedFile(filename, file.GoImportPath)

	// Build the model for every temporal service in the file up front so we can
	// detect cross-service identifier collisions before emitting anything.
	var services []*model.Service
	for _, s := range file.Services {
		if so, ok := proto.GetExtension(s.Desc.Options(), temporalv1.E_Service).(*temporalv1.ServiceOptions); !ok || so == nil {
			// not a temporal service if the `temporal.v1.service` option is not set
			continue
		}

		service, err := model.NewService(s, gen, config)
		if err != nil {
			return fmt.Errorf("failed to create service model: %w", err)
		}
		services = append(services, service)
	}

	// Detect package-level identifier collisions between services generated into
	// the same file. These identifiers are not service-prefixed, so two services
	// sharing a workflow/activity/signal/query/update method name would emit a
	// redeclaration and fail the consumer's build with no diagnostic here.
	if err := checkServiceCollisions(services, file.Proto.GetName()); err != nil {
		return err
	}

	// Build complete file content in memory first
	var fullContent bytes.Buffer
	fullContent.WriteString("// Code generated by protoc-gen-go-tmprl. DO NOT EDIT.\n")
	fullContent.WriteString("//\n")
	fullContent.WriteString("// version:\n")
	fullContent.WriteString("//   protoc-gen-go-tmprl version: " + version.Version + "\n")
	fullContent.WriteString("//   protoc-gen-go-tmprl commit: " + version.Commit + "\n")
	fullContent.WriteString("//\n")
	fullContent.WriteString("// source file: " + file.Proto.GetName() + "\n")
	fullContent.WriteString("\n")
	fullContent.WriteString("package " + string(file.GoPackageName) + "\n\n")

	for _, service := range services {
		// Create renderer
		r, err := renderer.NewRenderer(gen)
		if err != nil {
			return fmt.Errorf("failed to create renderer: %w", err)
		}

		// Render all components
		output, err := r.RenderAll(service)
		if err != nil {
			return fmt.Errorf("failed to render service %s: %w", service.GoName, err)
		}

		fullContent.WriteString(output)
	}

	// Format the complete file
	formatted, err := format.Source(fullContent.Bytes())
	if err != nil {
		return fmt.Errorf("failed to format generated code: %w", err)
	}

	// Write the formatted content
	gen.P(string(formatted))

	return nil
}

// checkServiceCollisions Returns an error if two temporal services in the same
// file emit the same package-level identifier. Such a collision compiles as
// valid Go syntax (so go/format accepts it) but redeclares an identifier and
// fails the consumer's build; catching it here fails loud at generation time.
func checkServiceCollisions(services []*model.Service, sourceFile string) error {
	owner := make(map[string]string)
	for _, s := range services {
		for _, name := range s.PackageScopedNames() {
			if prev, ok := owner[name]; ok {
				return fmt.Errorf("identifier %q is generated by both services %s and %s in %s; two services in one proto file cannot share a workflow/activity/signal/query/update method name", name, prev, s.GoName, sourceFile)
			}
			owner[name] = s.GoName
		}
	}
	return nil
}

func generateReadme(plugin *protogen.Plugin, file *protogen.File, config *model.Config) error {
	filename := file.GeneratedFilenamePrefix + "_tmprl_doc.md"

	needsGenerate := false
	for _, s := range file.Services {
		if so, ok := proto.GetExtension(s.Desc.Options(), temporalv1.E_Service).(*temporalv1.ServiceOptions); ok && so != nil {
			needsGenerate = true
		}
	}

	if !needsGenerate {
		return nil
	}

	// Use template-based renderer for documentation
	gen := plugin.NewGeneratedFile(filename, file.GoImportPath)
	gen.P(`<a id="top"></a>`)
	gen.P("# Services")

	for _, s := range file.Services {
		if so, ok := proto.GetExtension(s.Desc.Options(), temporalv1.E_Service).(*temporalv1.ServiceOptions); !ok || so == nil {
			// not a temporal service if the `temporal.v1.service` option is not set
			continue
		}

		// Create service model
		service, err := model.NewService(s, gen, config)
		if err != nil {
			return err
		}

		// Create renderer
		renderer, err := renderer.NewRenderer(gen)
		if err != nil {
			return err
		}

		// Render documentation
		doc, err := renderer.RenderDocumentation(service)
		if err != nil {
			return err
		}

		gen.P(doc)
	}

	gen.P("# Messages")

	// Create renderer for messages (reuse from last service)
	renderer, err := renderer.NewRenderer(gen)
	if err != nil {
		return err
	}

	for _, m := range file.Messages {
		msgDoc, err := renderer.RenderMessageDocumentation(m)
		if err != nil {
			return err
		}
		gen.P(msgDoc)
	}

	gen.P("\n\n[Back to top](#top)")

	return nil
}