From 7d2c4b8543ab8c1dc2bfe4fca5db07e83749775a Mon Sep 17 00:00:00 2001
From: bwireman <bwireman1@gmail.com>
Date: Wed, 9 Feb 2022 21:14:36 -0600
Subject: [PATCH] adds support for the gpb_header_include flag to correct
 imports

---
 README.md           | 10 ++++++---
 pkg/gleam/gpb.go    | 50 ++++++++++++++++++++++++++++++++++++++-------
 pkg/gleam/module.go | 28 ++++++++++++++++++-------
 3 files changed, 71 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index a5991cc..4ec4ea3 100644
--- a/README.md
+++ b/README.md
@@ -166,14 +166,18 @@ protoc --plugin=protoc-gen-gleam -I . --gleam_out="output_path=./src:./src" prot
 
 - 'output_path': (Required) specifies the desired output path
 - 'protoc_erl_path': path to gpb's protoc-erl
+- 'gpb_header_include': path to prepend to the header include for gpb. See [Issues](#Issues) for more info
+  - if you need a variable include here, remember that [erlang header resolution](https://www.erlang.org/doc/reference_manual/macros.html) is quite clever and can use environment variables
 
 ```bash
-protoc -I . --gleam_out="output_path=./src,protoc_erl_path=bin/protoc-erl:./src" protos/*.proto
+protoc -I . \
+  --gleam_out="gpb_header_include=$ENV/include/,output_path=./src,protoc_erl_path=bin/protoc-erl:./src" \
+  protos/*.proto
 ```
 
-### Issues
+### Known Issues
 
-You may need to manually update
+#### Includes aren't working?!
 
 ```erlang
 % generated in `gleam_gpb.erl`
diff --git a/pkg/gleam/gpb.go b/pkg/gleam/gpb.go
index 9aa619f..fc7492b 100644
--- a/pkg/gleam/gpb.go
+++ b/pkg/gleam/gpb.go
@@ -3,26 +3,31 @@ package gleam
 import (
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"os"
 	"os/exec"
+	"path"
+	"strings"
 )
 
 type gpbWrapper struct {
 	pathToBinary string
+	outputPath   string
 }
 
-func newGPBWrapper(pathToBinary string) (*gpbWrapper, error) {
-	if !exists(pathToBinary) {
+func newGPBWrapper(pathToBinary string, outputPath string) (*gpbWrapper, error) {
+	if _, exists := exists(pathToBinary); !exists {
 		return nil, fmt.Errorf("protoc-erl could not be found at %s", pathToBinary)
 	}
 
 	return &gpbWrapper{
 		pathToBinary: pathToBinary,
+		outputPath:   outputPath,
 	}, nil
 }
 
-func (g *gpbWrapper) generate(targets []string, outputPath string) (err error) {
-	args := []string{"-pkgs", "-modname", "gleam_gpb", "-I", ".", "-o", outputPath}
+func (g *gpbWrapper) generate(targets []string) (err error) {
+	args := []string{"-pkgs", "-modname", "gleam_gpb", "-I", ".", "-o", g.outputPath}
 	args = append(args, targets...)
 
 	cmd := exec.Command(g.pathToBinary, args...)
@@ -35,7 +40,38 @@ func (g *gpbWrapper) generate(targets []string, outputPath string) (err error) {
 	return nil
 }
 
-func exists(path string) bool {
-	_, err := os.Open(path)
-	return !errors.Is(err, os.ErrNotExist)
+const headerFileName = "gpb.hrl"
+
+func (g *gpbWrapper) updateImport(correctPath string) (err error) {
+	expectedPath := path.Join(g.outputPath, "gleam_gpb.erl")
+
+	if f, exists := exists(expectedPath); !exists {
+		return fmt.Errorf("updateImport failed: Could not find %s", expectedPath)
+	} else if b, err := ioutil.ReadAll(f); err != nil {
+		return err
+	} else {
+		contents := string(b)
+
+		newHeaderPath := path.Join(correctPath, headerFileName)
+		newContents := strings.Replace(contents, formatReplace(headerFileName), formatReplace(newHeaderPath), 1)
+
+		if err := os.WriteFile(expectedPath, []byte(newContents), 0666); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func formatReplace(path string) string {
+	return fmt.Sprintf("-include(\"%s\").", path)
+}
+
+func exists(path string) (*os.File, bool) {
+	f, err := os.Open(path)
+	if errors.Is(err, os.ErrNotExist) || err != nil {
+		return nil, false
+	}
+
+	return f, true
 }
diff --git a/pkg/gleam/module.go b/pkg/gleam/module.go
index 530e16a..6a469b5 100644
--- a/pkg/gleam/module.go
+++ b/pkg/gleam/module.go
@@ -10,23 +10,35 @@ import (
 
 type GleamModule struct {
 	*pgs.ModuleBase
-	tpl           *template.Template
-	protocErlPath string
-	output        string
+	tpl     *template.Template
+	wrapper *gpbWrapper
+
+	gpbHeaderInclude string
+	output           string
 }
 
 func Gleam() *GleamModule { return &GleamModule{ModuleBase: &pgs.ModuleBase{}} }
 
 func (g *GleamModule) InitContext(c pgs.BuildContext) {
 	g.ModuleBase.InitContext(c)
-	g.protocErlPath = g.Parameters().StrDefault("protoc_erl_path", "./deps/gpb/bin/protoc-erl")
 	g.tpl = template.Must(template.New("gleam-package-template").Parse(fields.GleamTemplate))
 
 	if g.output = g.Parameters().Str("output_path"); g.output == "" {
 		g.Fail("please specify the `output_path` flag")
 	}
+
+	protocErlPath := g.Parameters().StrDefault("protoc_erl_path", "./deps/gpb/bin/protoc-erl")
+	if wrapper, err := newGPBWrapper(protocErlPath, g.OutputPath()); err != nil {
+		g.Fail(err.Error())
+	} else {
+		g.wrapper = wrapper
+	}
+
+	g.gpbHeaderInclude = g.Parameters().Str("gpb_header_include")
 }
 
+func (g *GleamModule) OutputPath() string { return g.output }
+
 func (g *GleamModule) Name() string { return "gleam" }
 
 func (g *GleamModule) Execute(targets map[string]pgs.File, pkgs map[string]pgs.Package) []pgs.Artifact {
@@ -37,12 +49,14 @@ func (g *GleamModule) Execute(targets map[string]pgs.File, pkgs map[string]pgs.P
 		}
 	}
 
-	wrapper, err := newGPBWrapper(g.protocErlPath)
-	if err != nil {
+	if err := g.wrapper.generate(protosFilePaths); err != nil {
 		g.Fail(err.Error())
+	} else if g.gpbHeaderInclude != "" {
+		if err = g.wrapper.updateImport(g.gpbHeaderInclude); err != nil {
+			g.Fail(err.Error())
+		}
 	}
 
-	wrapper.generate(protosFilePaths, g.OutputPath())
 	g.AddCustomFile(g.OutputPath()+"/gleam_pb.gleam", fields.GleamPB, 0644)
 
 	for _, p := range pkgs {