main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"github.com/andrskom/jrpc2hh/gen/imports"
	"github.com/andrskom/jrpc2hh/gen/method"
	"github.com/andrskom/jrpc2hh/gen/service"
	"github.com/andrskom/jrpc2hh/gen/templates"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"strings"
	"text/template"
)

var hDir string
var pack string

func main() {
	flag.StringVar(&hDir, "s", "./testservice", "Service dir")
	flag.Parse()

	err := cleanAutoGeneratedFiles(hDir)
	logFatal("Error in autogenrated file deliting", err)

	fs := token.NewFileSet()
	packages, err := parser.ParseDir(fs, hDir, nil, parser.ParseComments)
	logFatal("Parsing dir error", err)

	regExpService, err := regexp.Compile("//[ ]*jrpc2hh:service")
	logFatal("Compiling regexp for service error", err)

	regExpMethod, err := regexp.Compile("//[ ]*jrpc2hh:method")
	logFatal("Compiling regep for method error", err)

	regExpMethodWithContext, err := regexp.Compile("//[ ]*jrpc2hh:method:withContext")
	logFatal("Compiling regep for method with context error", err)

	if len(packages) != 1 {
		log.Fatal("Expected that only one package will be parse")
	}
	for p, _ := range packages {
		pack = p
	}

	iMap, sl, ml := parse(regExpService, regExpMethod, regExpMethodWithContext, packages)
	iMap.GenerateAlias()
	generate(iMap, sl, ml)
}

func generate(iMap *imports.ImportMap, sl service.ServiceList, ml method.MethodList) {
	sTmpl, err := template.New("serviceTemplate").Parse(templates.Service)
	logFatal("Can't parse service template", err)

	mTmpl, err := template.New("methodTemplate").Parse(templates.Method)
	logFatal("Can't parse method template", err)

	for sn, sm := range ml {
		usedImports := make(map[string]string)
		usedImports["encoding/json"] = "json"
		usedImports["github.com/andrskom/jrpc2hh/models"] = "jModels"
		methods := make([]string, 0)
		for _, m := range sm {
			args := generateArgsBlock(m.Args, &usedImports, iMap)
			if m.ArgsWithContext {
				args = args + `
		args.WithContext(r.Context())`
			}
			res := generateResBlock(m.Result, &usedImports, iMap)

			buf := bytes.NewBuffer(make([]byte, 0))
			mTmpl.Execute(buf, struct {
				Method      string
				ArgsBlock   string
				ResultBlock string
			}{m.Name, args, res})
			methods = append(methods, buf.String())
		}

		file, err := os.OpenFile(fmt.Sprintf("%s/jrpc2hh_%s.go", hDir, strings.ToLower(sn)),
			os.O_WRONLY|os.O_CREATE|os.O_TRUNC,
			0755)
		if err != nil {
			log.Fatalf(fmt.Sprintf("Can't open file for writing generated data, %s", err.Error()))
		}
		sTmpl.Execute(file, struct {
			Package string
			Imports map[string]string
			Service string
			Methods []string
		}{pack, usedImports, sn, methods})
	}
}

func generateResBlock(res *method.Struct, ui *map[string]string, iMap *imports.ImportMap) string {
	if res.Pack+res.Name == "github.com/andrskom/jrpc2hh/modelsNilResult" {
		return `var res jModels.NilResult`
	} else {
		var resType string
		if res.Pack != "" {
			(*ui)[res.Pack] = iMap.GetFormattedAlias(res.Pack)
			resType = res.Prefix + res.Pack + "." + res.Name
		} else {
			resType = res.Prefix + res.Name
		}
		return "var res " + resType
	}
}

func generateArgsBlock(args *method.Struct, ui *map[string]string, iMap *imports.ImportMap) string {
	if args.Pack+args.Name == "github.com/andrskom/jrpc2hh/modelsNilArgs" {
		(*ui)[args.Pack] = iMap.GetFormattedAlias(args.Pack)
		return templates.ArgsEmpty
	} else {
		var argsType string
		if args.Pack != "" {
			(*ui)[args.Pack] = iMap.GetFormattedAlias(args.Pack)
			argsType = (*ui)[args.Pack] + "." + args.Name
		} else {
			argsType = args.Name
		}
		return fmt.Sprintf(templates.Args, argsType)
	}
}

func parse(regExpService *regexp.Regexp, regExpMethod *regexp.Regexp, regExpMethodWithContext *regexp.Regexp, packages map[string]*ast.Package) (*imports.ImportMap, service.ServiceList, method.MethodList) {
	iMap := imports.NewImportMap()
	sl := make(service.ServiceList)
	ml := make(method.MethodList)

	for _, p := range packages {
		for _, f := range p.Files {
			// collect imports
			localIMap := make(map[string]string)
			for _, im := range f.Imports {
				pV := strings.Trim(im.Path.Value, "\"")
				iMap.Register(pV)
				if im.Name != nil {
					localIMap[im.Name.String()] = pV
				} else {
					localIMap[pV] = pV
				}
			}

			// collect structs
			for _, d := range f.Decls {
				//collect services
				if reflect.TypeOf(d).Elem().Name() == "GenDecl" {
					gd, ok := (d).(*ast.GenDecl)
					if !ok {
						log.Fatal("Bad assertation type GenDecl")
					}
					if gd.Tok.String() == "type" && docHasMatch(regExpService, gd.Doc) {
						if len(gd.Specs) != 1 {
							log.Fatal("Bad Specs for type")
						}
						spec, ok := gd.Specs[0].(*ast.TypeSpec)
						if !ok {
							log.Fatal("Bad assertation TypeSpec")
						}
						if spec.Name == nil {
							log.Fatal("Bad Ident for types spec")
						}
						sl.Add(spec.Name.Name)
					}
					// collect methods
				} else if reflect.TypeOf(d).Elem().Name() == "FuncDecl" {
					fd, ok := (d).(*ast.FuncDecl)
					if !ok {
						log.Fatal("Bad assertation func FunDecl")
					}
					if docHasMatch(regExpMethod, fd.Doc) {
						mN := fd.Name.String()
						if fd.Recv == nil {
							log.Fatal("Recv of service method can't be nil")
						}
						if len(fd.Recv.List) != 1 {
							log.Fatal("List of recv service method can't be nil")
						}
						mT, ok := (fd.Recv.List[0].Type).(*ast.StarExpr)
						if !ok {
							log.Fatal("Type associated with method must be pointer")
						}
						i, ok := (mT.X).(*ast.Ident)
						if !ok {
							log.Fatal("Bad Ident for method spec")
						}
						assType := i.Name

						if len(fd.Type.Params.List) != 2 {
							log.Fatal("Count of params must be equal 2")
						}
						argsAst := fd.Type.Params.List[0]
						resAst := fd.Type.Params.List[1]

						var args *method.Struct
						var res *method.Struct

						switch reflect.TypeOf(argsAst.Type).String() {
						case "*ast.SelectorExpr":
							t := (argsAst.Type).(*ast.SelectorExpr)
							x, ok := (t.X).(*ast.Ident)
							if !ok {
								log.Fatal("Bad assertation type")
							}
							if _, ok := localIMap[x.Name]; !ok {
								log.Fatal("Problem with inport and alias")
							}
							args = method.NewStruct(localIMap[x.Name], t.Sel.Name)
						case "*ast.Ident":
							t := (argsAst.Type).(*ast.Ident)
							args = method.NewStruct("", t.Name)
						default:
							log.Fatal("Unknown type of args")
						}

						if reflect.TypeOf(resAst.Type).String() != "*ast.StarExpr" {
							log.Fatal("Result must be pointer")
						}

						sE, _ := (resAst.Type).(*ast.StarExpr)

						switch reflect.TypeOf(sE.X).String() {
						case "*ast.StarExpr":
							t := (sE.X).(*ast.StarExpr)
							switch reflect.TypeOf(t.X).String() {
							case "*ast.Ident":
								t := (t.X).(*ast.Ident)
								res = method.NewStruct("", t.Name)
								res.SetPrefix("*")
							case "*ast.SelectorExpr":
								t := (t.X).(*ast.SelectorExpr)
								x, ok := (t.X).(*ast.Ident)
								if !ok {
									log.Fatal("Bad assertation type")
								}
								if _, ok := localIMap[x.Name]; !ok {
									log.Fatal("Problem with inport and alias")
								}
								res = method.NewStruct(localIMap[x.Name], t.Sel.Name)
								res.SetPrefix("*")
							default:
								log.Fatal("Unknown type of res")
							}
						case "*ast.SelectorExpr":
							t := (sE.X).(*ast.SelectorExpr)
							x, ok := (t.X).(*ast.Ident)
							if !ok {
								log.Fatal("Bad assertation type")
							}
							if _, ok := localIMap[x.Name]; !ok {
								log.Fatal("Problem with inport and alias")
							}
							res = method.NewStruct(localIMap[x.Name], t.Sel.Name)
						case "*ast.Ident":
							t := (sE.X).(*ast.Ident)
							res = method.NewStruct("", t.Name)
						default:
							log.Fatal("Unknown type of res")
						}
						withContext := docHasMatch(regExpMethodWithContext, fd.Doc)
						ml.Add(assType, method.NewMethod(mN, args, res, withContext))
					}
				}
			}
		}
	}

	return iMap, sl, ml
}

func docHasMatch(regexp *regexp.Regexp, doc *ast.CommentGroup) bool {
	res := false
	if doc != nil {
		for _, cm := range doc.List {
			if regexp.Match([]byte(cm.Text)) {
				res = true
			}
		}
	}

	return res
}

func logFatal(comment string, err error) {
	if err != nil {
		log.Fatal(fmt.Sprintf("%s: %s", comment, err.Error()))
	}
}

func cleanAutoGeneratedFiles(sDir string) error {
	regExpFile, err := regexp.Compile(`.*/jrpc2hh_.+\.go`)
	if err != nil {
		return err
	}
	dFun := func(path string, f os.FileInfo, err error) error {
		if regExpFile.Match([]byte(path)) {
			os.Remove(path)
		}
		return nil
	}
	return filepath.Walk(sDir, dFun)
}