package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/fs"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"golang.org/x/mod/modfile"
)

var (
	file     = flag.String("mod", "go.mod", "go.mod file")
	cache    = flag.String("modcache", getmodcache(), "location of go mod cache")
	excheckm = flag.String("exempt", "golang.org", "domains exempt from the search, seperated by space")

	checkInit   = flag.Bool("ci", false, "check for implementations of init()")
	checkUnsafe = flag.Bool("cu", false, "check for imports of unsafe")
	checkCgo    = flag.Bool("cc", false, "check for imports of cgo")
	checkPprof  = flag.Bool("cp", false, "check for imports of net/http/pprof")

	exceptions []string

	relcache string
)

func getmodcache() string {
	if c := os.Getenv("GOMODCACHE"); c != "" {
		return c
	} else if c = os.Getenv("GOPATH"); c != "" {
		return c + "/pkg/mod"
	} else {
		return os.Getenv("HOME") + "/pkg/mod"
	}

}

func main() {
	flag.Parse()

	exceptions = strings.Fields(*excheckm)
	relcache = strings.Replace(*cache, os.Getenv("HOME"), "~", 1)

	if !(*checkInit || *checkUnsafe || *checkCgo || *checkPprof) {
		fmt.Println("Nothing to check. Use -ci|-cu|-cc|-cp")
		return
	}

	data, err := ioutil.ReadFile(*file)
	if err != nil {
		log.Fatal(err)
	}

	mf, err := modfile.Parse(*file, data, nil)
	if err != nil {
		log.Fatal(err)
	}

	var wg sync.WaitGroup
	wg.Add(len(mf.Require))
	for _, r := range mf.Require {
		path := strings.Join(r.Syntax.Token, "@")
		go func() {
			defer wg.Done()
			checkPath(*cache + "/" + path)
		}()
	}

	wg.Wait()

}

func isExempt(path string) bool {
	for _, pattern := range exceptions {
		if strings.Contains(path, pattern) {
			return true
		}
	}
	return false
}

func modpath(filename string) string {
	return filepath.Join(relcache, strings.TrimPrefix(filename, *cache+"/"))
}

func checkPath(path string) {
	var fset = token.NewFileSet()

	filter := func(inf fs.FileInfo) bool {
		return !strings.HasSuffix(inf.Name(), "_test.go")
	}

	filepath.WalkDir(path, func(p string, info fs.DirEntry, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			return nil
		} else if isExempt(p) {
			return nil
		}

		pkgs, err := parser.ParseDir(fset, p, filter, parser.Mode(0))
		if err != nil {
			return err
		}

		for _, pkg := range pkgs {
			for filename, file := range pkg.Files {
				if *checkUnsafe || *checkCgo || *checkPprof {
					for _, imp := range file.Imports {
						if *checkUnsafe && imp.Path.Value == `"unsafe"` {
							fmt.Println("unsafe in", modpath(filename))
						} else if *checkCgo && imp.Path.Value == `"C"` {
							fmt.Println("cgo    in", modpath(filename))
						} else if *checkPprof && imp.Path.Value == `"net/http/pprof"` {
							fmt.Println("pprof  in", modpath(filename))
						}
					}
				}

				if !*checkInit {
					break
				}
			INIT:
				for _, decl := range file.Decls {
					if f, ok := decl.(*ast.FuncDecl); ok {
						if f.Name.Name == "init" {
							fmt.Println("init   in", modpath(filename))
							break INIT
						}
					}
				}
			}
		}

		return nil
	})
}