//+build ignore

package main

import (
	"flag"
	"fmt"
	"go/build"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
)

const (
	thirdPartyDir = "code.justin.tv/common/chitin/third_party"
)

var (
	verbose = flag.Bool("verbose", false, "Print additional information about what the tool is doing")
)

func verboseLogf(format string, a ...interface{}) {
	if *verbose {
		log.Printf("[chitin/protoc_gen] %s", fmt.Sprintf(format, a...))
	}
}

func protocCommand(rel string, includes, protoFiles []string) *exec.Cmd {
	// protoc -I.:$GOPATH/src:/usr/local/include --go_out=plugins=grpc:. event.proto
	out := ""
	args := []string{"protoc"}
	for _, include := range includes {
		args = append(args, "--proto_path", include)
		if rel == include || strings.HasPrefix(rel, include+"/") {
			if out == "" {
				out = include
			} else if len(include) > len(out) {
				log.Fatalf(`FATAL:

The directory we're about to build shows up in the include path more than
once, and the first occurrence is not the longest. The protobuf compiler will
complain about some identifiers already being defined.

Be sure that any imports with common prefixes are ordered with the most
specific directories first.

Target directory:                %s
First matching import path:      %s
Subsequent matching import path: %s

`, rel, out, include)
			}
		}
	}

	// Options of protoc-gen-go, including for import maps, are documented here:
	//
	// https://github.com/golang/protobuf/blob/master/README.md#parameters

	plugins := []string{"grpc"}

	importMaps := map[string]string{
		"google/protobuf/descriptor.proto": "github.com/golang/protobuf/protoc-gen-go/descriptor",
		"google/api/annotations.proto":     "google.golang.org/genproto/googleapis/api/annotations",
		"google/protobuf/field_mask.proto": "google.golang.org/genproto/protobuf/field_mask",
	}
	var importMapOptions []string
	for k, v := range importMaps {
		importMapOptions = append(importMapOptions, fmt.Sprintf("M%s=%s", k, v))
	}
	sort.Strings(importMapOptions)

	var options []string
	options = append(options, importMapOptions...)
	if len(plugins) > 0 {
		options = append(options, "plugins="+strings.Join(plugins, "+"))
	}

	args = append(args, "--go_out", strings.Join(options, ",")+":"+out)
	for _, file := range protoFiles {
		args = append(args, filepath.Join(rel, file))
	}
	verboseLogf("exec: %v", args)
	cmd := exec.Command(args[0], args[1:]...)
	return cmd
}

func main() {
	var (
		startDir = flag.String("start", ".", "Walk the filesystem starting at the specified directory")
	)
	flag.Parse()
	log.SetFlags(0)

	err := os.Chdir(*startDir)
	if err != nil {
		log.Fatalf("could not change directory: %v", err)
	}

	pwd, err := os.Getwd()
	if err != nil {
		log.Fatalf("could not get working directory: %v", err)
	}

	ctxt := new(build.Context)
	*ctxt = build.Default

	includes := existantDirs(findIncludes(ctxt, pwd))
	verboseLogf("includes= %q", includes)

	failed := false

	files, err := findProtoFiles(".")
	if err != nil {
		log.Fatalf("could not find proto files: %v", err)
	}
	verboseLogf("findProtoFiles() returned files= %q", files)

	var dirs []string
	for dir := range files {
		if dir == "third_party/google/protobuf" {
			// The canonical go files for these live in the
			// github.com/golang/protobuf, in
			// github.com/golang/protobuf/ptypes and in
			// github.com/golang/protobuf/protoc-gen-go/descriptor
			continue
		}
		dirs = append(dirs, dir)
	}
	sort.Strings(dirs)

	for _, dir := range dirs {
		full := filepath.Join(pwd, dir)
		protoFiles := files[dir]
		verboseLogf("main loop: dir=%q protoFiles=%q", dir, protoFiles)

		cmd := protocCommand(full, includes, protoFiles)
		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
		err = cmd.Run()
		if err != nil {
			log.Printf("protoc failed on dir %q", dir)
			failed = true
			continue
		}
	}

	for _, dir := range dirs {
		full := filepath.Join(pwd, dir)

		err := applyEgTemplates(ctxt, full)
		if err != nil {
			log.Printf("%v", err)
			failed = true
			if isFatal(err) {
				os.Exit(1)
			}
			continue
		}
	}

	if failed {
		os.Exit(1)
	}
}

func applyEgTemplates(ctxt *build.Context, full string) error {
	pkg, err := ctxt.ImportDir(full, 0)
	if err != nil {
		return fmt.Errorf("could not load package in dir %q: %v", full, err)
	}
	for _, tmpl := range []string{
		egInvoke,
		egNewClientStream,
		egRegisterService,
	} {
		args := []string{"eg", "-w", "-t", "/dev/stdin", "--"}
		args = append(args, pkg.GoFiles...)
		verboseLogf("exec: %v", args)
		cmd := exec.Command(args[0], args[1:]...)
		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
		cmd.Stdin = strings.NewReader(tmpl)
		cmd.Dir = full
		err = cmd.Run()
		if err != nil {
			if ee, ok := err.(*exec.Error); ok && ee.Err == exec.ErrNotFound {
				return fatal(fmt.Errorf("" +
					"The eg tool was not found in your PATH. You can install it via go get:\n" +
					"    $ go get golang.org/x/tools/cmd/eg"))
			}
			return fmt.Errorf("eg failed on dir %q: %v", full, err)
		}
	}

	return nil
}

func repoRoot(pwd string) (root string, err error) {
	for dir := pwd; ; dir = filepath.Dir(dir) {
		for _, entry := range []string{".git", "vendor", "Godeps"} {
			_, err := os.Stat(filepath.Join(dir, entry))
			if err == nil {
				return dir, nil
			}
			if !os.IsNotExist(err) {
				return "", err
			}
		}
		if dir == filepath.Dir(dir) {
			// We've reached the filesystem root without finding the repo root
			return "", os.ErrNotExist
		}
	}
}

func findProtoFiles(dir string) (map[string][]string, error) {
	files := make(map[string][]string)

	err := filepath.Walk(dir, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() && (info.Name() == "vendor" || info.Name() == "Godeps") {
			return filepath.SkipDir
		}
		if !info.Mode().IsRegular() || !strings.HasSuffix(info.Name(), ".proto") {
			return nil
		}
		dir, base := filepath.Dir(name), filepath.Base(name)
		files[dir] = append(files[dir], base)
		return nil
	})
	if err != nil {
		return nil, err
	}

	return files, nil
}

func findIncludes(ctxt *build.Context, pwd string) []string {
	var includes []string

	root, err := repoRoot(pwd)
	if err != nil {
		return includes
	}

	for _, dir := range []string{
		"./vendor",                // GO15VENDOREXPERIMENT
		"./Godeps/_workspace/src", // godep

		// This one is to reverse "code.justin.tv/org/repo" to find this
		// dir's GOPATH entry.
		"../../..",
	} {
		for _, incl := range []string{
			// This directory contains proto definitions from third party
			// developers, such as google/protobuf/descriptor.proto
			thirdPartyDir,

			// This allows importing proto files with the same
			// qualifications used for Go imports. For example, imports
			// starting with "code.justin.tv/web/thoth" will work just as
			// well whether the web/thoth repo is checked out in your
			// GOPATH or has been vendored into your repo via godep.
			".",
		} {
			incl := filepath.Join(root, dir, incl)
			includes = append(includes, incl)
		}
	}

	return includes
}

func existantDirs(incls []string) []string {
	var includes []string
	for _, incl := range incls {
		if stat, err := os.Stat(incl); err == nil && stat.IsDir() {
			includes = append(includes, incl)
		}
	}
	return includes
}

type augmentedError struct {
	err   error
	fatal bool
}

func (ae *augmentedError) Err() error {
	return ae.err
}

func (ae *augmentedError) Error() string {
	return ae.err.Error()
}

func fatal(err error) error {
	return &augmentedError{fatal: true, err: err}
}

func isFatal(err error) bool {
	if ae, ok := err.(*augmentedError); ok && ae.fatal {
		return true
	}
	return false
}

var (
	egInvoke = `
package eg

import (
	"code.justin.tv/common/chitin/grpctrace"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

func before(ctx context.Context, method string, args, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
	return grpc.Invoke(ctx, method, args, reply, cc, opts...)
}

func after(ctx context.Context, method string, args, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
	return grpctrace.Invoke(ctx, method, args, reply, cc, opts...)
}
`[1:]

	egNewClientStream = `
package eg

import (
	"code.justin.tv/common/chitin/grpctrace"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

func before(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
	return grpc.NewClientStream(ctx, desc, cc, method, opts...)
}

func after(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
	return grpctrace.NewClientStream(ctx, desc, cc, method, opts...)
}
`[1:]

	egRegisterService = `
package eg

import (
	"code.justin.tv/common/chitin/grpctrace"
	"google.golang.org/grpc"
)

func before(s *grpc.Server, sd *grpc.ServiceDesc, ss interface{}) {
	s.RegisterService(sd, ss)
}

func after(s *grpc.Server, sd *grpc.ServiceDesc, ss interface{}) {
	grpctrace.RegisterService(s, sd, ss)
}
`[1:]
)
