/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"bytes"
	goflag "flag"
	"fmt"
	"go/build"
	"io"
	"os"
	"sort"
	"strings"

	"github.com/spf13/pflag"
)

var flPrune = pflag.StringSlice("prune", nil, "sub-packages to prune (recursive, may be specified multiple times)")
var flDebug = pflag.BoolP("debug", "d", false, "enable debugging output")
var flHelp = pflag.BoolP("help", "h", false, "print help and exit")

func main() {
	pflag.CommandLine.AddGoFlagSet(goflag.CommandLine)
	pflag.Usage = func() { help(os.Stderr) }
	pflag.Parse()

	debug("PWD", getwd())

	build.Default.BuildTags = []string{"ignore_autogenerated"}
	build.Default.UseAllFiles = false

	if *flHelp {
		help(os.Stdout)
		os.Exit(0)
	}
	if len(pflag.Args()) == 0 {
		help(os.Stderr)
		os.Exit(1)
	}
	for _, in := range pflag.Args() {
		if strings.HasSuffix(in, "/...") {
			// Recurse.
			debug("starting", in)
			pkgName := strings.TrimSuffix(in, "/...")
			if err := WalkPkg(pkgName, visitPkg); err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
		} else {
			// Import one package.
			if err := saveImport(in); err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(2)
			}
		}
	}
}

func help(out io.Writer) {
	fmt.Fprintf(out, "Usage: %s [FLAG...] <PKG...>\n", os.Args[0])
	fmt.Fprintf(out, "\n")
	fmt.Fprintf(out, "go2make calculates all of the dependencies of a set of Go packages and prints\n")
	fmt.Fprintf(out, "them as variable definitions suitable for use as a Makefile.\n")
	fmt.Fprintf(out, "\n")
	fmt.Fprintf(out, "Package specifications may be simple (e.g. 'example.com/txt/color') or\n")
	fmt.Fprintf(out, "recursive (e.g. 'example.com/txt/...')\n")
	fmt.Fprintf(out, " Example:\n")
	fmt.Fprintf(out, "  $ %s ./example.com/pretty\n", os.Args[0])
	fmt.Fprintf(out, "  example.com/txt/split := \\\n")
	fmt.Fprintf(out, "    /go/src/example.com/txt/split/ \\\n")
	fmt.Fprintf(out, "    /go/src/example.com/txt/split/split.go \\\n")
	fmt.Fprintf(out, "  ./example.com/pretty := \\\n")
	fmt.Fprintf(out, "    /go/src/example.com/pretty/ \\\n")
	fmt.Fprintf(out, "    /go/src/example.com/pretty/print.go \\\n")
	fmt.Fprintf(out, "    /go/src/example.com/txt/split/ \\\n")
	fmt.Fprintf(out, "    /go/src/example.com/txt/split/split.go\n")
	fmt.Fprintf(out, "\n")
	fmt.Fprintf(out, " Flags:\n")

	pflag.PrintDefaults()
}

func debug(items ...interface{}) {
	if *flDebug {
		x := []interface{}{"DBG:"}
		x = append(x, items...)
		fmt.Println(x...)
	}
}

func visitPkg(importPath, absPath string) error {
	debug("visit", importPath)
	return saveImport(importPath)
}

func prune(pkgName string) bool {
	for _, pr := range *flPrune {
		if pr == pkgName {
			return true
		}
	}
	return false
}

// cache keeps track of which packages we have already loaded.
var cache = map[string]*build.Package{}

func saveImport(pkgName string) error {
	if cache[pkgName] != nil {
		return nil
	}
	if prune(pkgName) {
		debug("prune", pkgName)
		return ErrSkipPkg
	}
	pkg, err := loadPackage(pkgName)
	if err != nil {
		return err
	}
	debug("save", pkgName)
	cache[pkgName] = pkg

	debug("recurse", pkgName)
	defer func() { debug("done   ", pkgName) }()
	if !pkg.Goroot && (len(pkg.GoFiles)+len(pkg.Imports) > 0) {
		// Process deps of this package before the package itself.
		for _, impName := range pkg.Imports {
			if impName == "C" {
				continue
			}
			debug("depends on", impName)
			saveImport(impName)
		}

		// Emit a variable for each package.
		var buf bytes.Buffer
		buf.WriteString(pkgName)
		buf.WriteString(" := ")

		// Packages depend on their own directories, their own files, and
		// transitive list of all deps' directories and files.
		all := map[string]struct{}{}
		all[pkg.Dir+"/"] = struct{}{}
		filesForPkg(pkg, all)
		for _, imp := range pkg.Imports {
			pkg := cache[imp]
			if pkg == nil || pkg.Goroot {
				continue
			}
			all[pkg.Dir+"/"] = struct{}{}
			filesForPkg(pkg, all)
		}
		// Sort and de-dup them.
		files := flatten(all)
		for _, f := range files {
			buf.WriteString(" \\\n  ")
			buf.WriteString(f)
		}

		fmt.Println(buf.String())
	}
	return nil
}

func filesForPkg(pkg *build.Package, all map[string]struct{}) {
	for _, file := range pkg.GoFiles {
		if pkg.Dir != "." {
			file = pkg.Dir + "/" + file
		}
		all[file] = struct{}{}
	}
}

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

func loadPackage(pkgName string) (*build.Package, error) {
	debug("load", pkgName)
	pkg, err := build.Import(pkgName, getwd(), 0)
	if err != nil {
		// We can ignore NoGoError.  Anything else is real.
		if _, ok := err.(*build.NoGoError); !ok {
			return nil, err
		}
	}
	return pkg, nil
}

func getwd() string {
	pwd, err := os.Getwd()
	if err != nil {
		panic(fmt.Sprintf("can't get working directory: %v", err))
	}
	return pwd
}