service.proto.go

package template

import (
	"context"
	"fmt"
	"sort"

	"github.com/devimteam/microgen/generator/strings"
	"github.com/devimteam/microgen/generator/write_strategy"
	"github.com/vetcher/go-astra/types"
)

const tab = "    "

type protoTemplate struct {
	info     *GenerationInfo
	protoPkg string
}

func NewProtoTemplate(info *GenerationInfo, protoPkg string) Template {
	return &protoTemplate{
		info:     info,
		protoPkg: protoPkg,
	}
}

// Renders stub code for service, its methods and constructor, that implements service interface.
//
//		// Generated by "microgen" tool.
//		// Structure stringService implements StringService interface.
//		type stringService struct {
//		}
//
//		func NewStringService() StringService {
//			panic("constructor not provided")
//		}
//
//		func (s *stringService) Count(ctx context.Context, text string, symbol string) (count int, positions []int) {
//			panic("method not provided")
//		}
//
func (t *protoTemplate) Render(ctx context.Context) write_strategy.Renderer {
	f := BufferAdapter{}
	f.Ln(`syntax = "proto3";`)
	f.Ln()
	f.Lnf(`option go_package = "%s;pb";`, t.info.ProtobufPackageImport)
	f.Ln()
	f.Lnf("package %s;", t.protoPkg)
	f.Ln()
	{
		d := f.Hold()
		imports := make(map[string]struct{})
		// Draw service
		d.Lnf("service %s {", t.info.Iface.Name)
		for _, method := range t.info.Iface.Methods {
			if !t.info.AllowedMethods[method.Name] {
				continue
			}
			reqTypeName, externalImport := protoMessageName(RemoveContextIfFirst(method.Args), requestStructName(method))
			if externalImport != nil {
				imports[*externalImport] = struct{}{}
			}
			respTypeName, externalImport := protoMessageName(removeErrorIfLast(method.Results), responseStructName(method))
			if externalImport != nil {
				imports[*externalImport] = struct{}{}
			}
			d.Lnf(tab+"rpc %s (%s) returns (%s);", method.Name, reqTypeName, respTypeName)
		}
		d.Ln("}")

		// Draw message types
		for _, method := range t.info.Iface.Methods {
			if !t.info.AllowedMethods[method.Name] {
				continue
			}
			{
				args := RemoveContextIfFirst(method.Args)
				reqTypeName, externalImport := protoMessageName(args, requestStructName(method))
				if externalImport == nil {
					d.Ln()
					d.Lnf("message %s {", reqTypeName)
					for i, arg := range args {
						n, imp := protoTypeName(arg.Type)
						if imp != nil {
							imports[*imp] = struct{}{}
						}
						d.Lnf(tab+"%s %s = %d;", n, strings.ToSnakeCase(arg.Name), i+1)
					}
					d.Ln("}")
				}
			}
			{
				params := removeErrorIfLast(method.Results)
				reqTypeName, externalImport := protoMessageName(params, responseStructName(method))
				if externalImport == nil {
					d.Ln()
					d.Lnf("message %s {", reqTypeName)
					for i, arg := range params {
						n, imp := protoTypeName(arg.Type)
						if imp != nil {
							imports[*imp] = struct{}{}
						}
						d.Lnf(tab+"%s %s = %d;", n, strings.ToSnakeCase(arg.Name), i+1)
					}
					d.Ln("}")
				}
			}
		}

		for _, imp := range sortedSliceFromStringSet(imports) {
			f.Lnf(`import "%s";`, imp)
		}

		f.Ln()
		d.Release()
	}
	return f
}

func (protoTemplate) DefaultPath() string {
	return "service.proto"
}

func (t *protoTemplate) Prepare(ctx context.Context) error {
	return nil
}

func (t *protoTemplate) ChooseStrategy(ctx context.Context) (write_strategy.Strategy, error) {
	return write_strategy.NewCreateRawFileStrategy(t.info.OutputFilePath, t.DefaultPath()), nil
}

const (
	googleProtobuf             = "google.protobuf."
	googleProtobufEmpty        = googleProtobuf + "Empty"
	googleProtobufStringValue  = googleProtobuf + "StringValue"
	googleProtobufBoolValue    = googleProtobuf + "BoolValue"
	googleProtobufInt64Value   = googleProtobuf + "Int64Value"
	googleProtobufUInt64Value  = googleProtobuf + "UInt64Value"
	googleProtobufInt32Value   = googleProtobuf + "Int32Value"
	googleProtobufUInt32Value  = googleProtobuf + "UInt32Value"
	googleProtobufFloat64Value = googleProtobuf + "DoubleValue"
	googleProtobufFloat32Value = googleProtobuf + "FloatValue"
	googleProtobufTimestamp    = googleProtobuf + "Timestamp"

	importGoogleProtobuf          = "google/protobuf/"
	importGoogleProtobufWrappers  = importGoogleProtobuf + "wrappers.proto"
	importGoogleProtobufEmpty     = importGoogleProtobuf + "empty.proto"
	importGoogleProtobufTimestamp = importGoogleProtobuf + "timestamp.proto"
)

func protoMessageName(params []types.Variable, def string) (string, *string) {
	switch len(params) {
	case 0:
		return googleProtobufEmpty, sp(importGoogleProtobufEmpty)
	case 1:
		switch params[0].Type.String() {
		case "*string":
			return googleProtobufStringValue, sp(importGoogleProtobufWrappers)
		case "*bool":
			return googleProtobufBoolValue, sp(importGoogleProtobufWrappers)
		case "*int64":
			return googleProtobufInt64Value, sp(importGoogleProtobufWrappers)
		case "*int32":
			return googleProtobufInt32Value, sp(importGoogleProtobufWrappers)
		case "*uint64":
			return googleProtobufUInt64Value, sp(importGoogleProtobufWrappers)
		case "*uint32":
			return googleProtobufUInt32Value, sp(importGoogleProtobufWrappers)
		case "*float64":
			return googleProtobufFloat64Value, sp(importGoogleProtobufWrappers)
		case "*float32":
			return googleProtobufFloat32Value, sp(importGoogleProtobufWrappers)
		case "time.Time", "*time.Time":
			return googleProtobufTimestamp, sp(importGoogleProtobufTimestamp)
		}
		fallthrough
	default:
		return def, nil
	}
}

func protoTypeName(v types.Type) (t string, imp *string) {
	switch v.String() {
	case "*string":
		return googleProtobufStringValue, sp(importGoogleProtobufWrappers)
	case "*bool":
		return googleProtobufBoolValue, sp(importGoogleProtobufWrappers)
	case "*int64":
		return googleProtobufInt64Value, sp(importGoogleProtobufWrappers)
	case "*int32":
		return googleProtobufInt32Value, sp(importGoogleProtobufWrappers)
	case "*uint64":
		return googleProtobufUInt64Value, sp(importGoogleProtobufWrappers)
	case "*uint32":
		return googleProtobufUInt32Value, sp(importGoogleProtobufWrappers)
	case "*float64":
		return googleProtobufFloat64Value, sp(importGoogleProtobufWrappers)
	case "*float32":
		return googleProtobufFloat32Value, sp(importGoogleProtobufWrappers)
	case "time.Time", "*time.Time":
		return googleProtobufTimestamp, sp(importGoogleProtobufTimestamp)
	case "[]byte":
		t = "bytes"
	case "int":
		t = "int64"
	case "uint":
		t = "uint64"
	default:
		if n := types.TypeName(v); n != nil {
			t = *n
		}
	}
	if types.IsMap(v) {
		m := types.TypeMap(v).(types.TMap)
		key, _ := protoTypeName(m.Key)
		value, _ := protoTypeName(m.Value)
		t = fmt.Sprintf("map<%s, %s>", key, value)
	}
	if types.IsArray(v) {
		v, _ := protoTypeName(types.TypeArray(v).(types.LinearType).NextType())
		t = "repeated " + v
	}
	return t, nil
}

func sp(s string) *string {
	return &s
}

func sortedSliceFromStringSet(s map[string]struct{}) []string {
	slice := make([]string, 0, len(s))
	for k := range s {
		slice = append(slice, k)
	}
	sort.Strings(slice)
	return slice
}