package main

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

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

var (
	dir      = flag.String("d", ".", "directory with go.mod file")
	recurse  = flag.Bool("r", false, "recursively search for go.mod files")
	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"
	}

}

type pathcache struct {
	m sync.RWMutex
	c map[string]bool
}

func (p *pathcache) get(path string) (ok bool) {
	p.m.RLock()
	defer p.m.RUnlock()
	_, ok = p.c[path]
	return
}

func (p *pathcache) add(path string) {
	p.m.Lock()
	defer p.m.Unlock()
	p.c[path] = true
}

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
	}

	pc := &pathcache{
		c: map[string]bool{},
	}
	if !*recurse {
		file := filepath.Join(*dir, "go.mod")
		singleFile(file, pc)
	} else {
		recursive(*dir, pc)
	}
}

func singleFile(file string, pc *pathcache) {
	data, err := os.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
	for _, r := range mf.Require {
		path := strings.Join(r.Syntax.Token, "@")
		if done := pc.get(path); done {
			continue
		}
		pc.add(path)
		wg.Add(1)
		go func() {
			defer wg.Done()
			checkPath(*cache + "/" + path)
		}()
	}

	wg.Wait()
}

func recursive(dir string, pc *pathcache) {
	filepath.WalkDir(dir, func(p string, info fs.DirEntry, err error) error {
		if err != nil {
			return err
		} else if info.IsDir() {
			return nil
		} else if "go.mod" == filepath.Base(p) {
			fmt.Printf("analyzing %s\n", p)
			singleFile(p, pc)
		}
		return nil
	})
}

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
	})
}