//go:build linux

package main

import (
	"debug/elf"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

var (
	destFlag    string
	packageFlag string
)

func main() {
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, `Usage: memexec-gen PATH

Environment variables:
  LD_LIBRARY_PATH
  	list of directories where to search libraries, see 'man 7 ld.so'

Options:
`)
		flag.PrintDefaults()
	}
	flag.StringVar(&destFlag, "dest", "",
		"`path` where binaries and go file are written to, defaults to PATH basename",
	)
	flag.StringVar(&packageFlag, "package", "",
		"generated package `name`, defaults to PATH basename",
	)
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(2)
	}
	if err := run(flag.Arg(0)); err != nil {
		fmt.Fprintf(os.Stderr, "error: %s\n", err)
		os.Exit(127)
	}
}

func run(bin string) error {
	base := filepath.Base(bin)
	if destFlag == "" {
		destFlag = base
	}
	if packageFlag == "" {
		packageFlag = base
	}

	err := os.MkdirAll(destFlag, 0o755)
	if err != nil && !os.IsExist(err) {
		return err
	}
	output, err = os.OpenFile(
		filepath.Join(destFlag, "memexec.go"),
		os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
		0o644,
	)
	if err != nil {
		return err
	}

	libs := map[string]string{}
	if err := ldd(libs, bin, func(name string) (string, error) {
		return name, nil
	}); err != nil {
		return err
	}
	for name, path := range libs {
		src, err := os.Open(path)
		if err != nil {
			return err
		}
		if name == bin {
			name = base
		}
		dst, err := os.OpenFile(
			filepath.Join(destFlag, name),
			os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
			0o755,
		)
		if err != nil {
			return err
		}
		if _, err := io.Copy(dst, src); err != nil {
			return err
		}
		if err := dst.Close(); err != nil {
			return err
		}
		if err := src.Close(); err != nil {
			return err
		}
	}

	out("// Code generated by memexec-gen. DO NOT EDIT.")
	out("package ", packageFlag)
	out()
	out("import (")
	out(`	_ "embed"`)
	out(`	"os"`)
	out(`	"os/exec"`)
	out(`	"path/filepath"`)
	out()
	out(`	"github.com/amenzhinsky/go-memexec"`)
	out(")")
	out()

	for name := range libs {
		if name == bin {
			name = base
		}
		out("//go:embed ", name)
		out("var ", goName(name), " []byte")
		out()
	}

	out("func New(opts ...memexec.Option) (*memexec.Exec, error) {")
	out(`	temp, err := os.MkdirTemp("", "")`)
	out("	if err != nil {")
	out("		return nil, err")
	out("	}")
	for name := range libs {
		if name == bin {
			continue
		}
		out(`	if err := os.WriteFile(filepath.Join(temp, "`, name, `"), `, goName(name), ", 0o755); err != nil {")
		out("		return nil, err")
		out("	}")
	}
	out("	return memexec.New(", goName(base), ",  append([]memexec.Option{")
	out("		memexec.WithPrepare(func(exe *exec.Cmd) {")
	out(`			exe.Env = append(exe.Env, "LD_LIBRARY_PATH="+temp)`)
	out("		}),")
	out("		memexec.WithCleanup(func() error {")
	out("			return os.RemoveAll(temp)")
	out("		}),")
	out("	}, opts...)...)")
	out("}")

	return nil
}

var output io.Writer

func out(strs ...string) {
	for _, str := range strs {
		if _, err := fmt.Fprint(output, str); err != nil {
			panic(err)
		}
	}
	if _, err := fmt.Fprintln(output); err != nil {
		panic(err)
	}
}

func goName(name string) string {
	name = strings.ReplaceAll(name, ".", "_")
	name = strings.ReplaceAll(name, "-", "_")
	return "memexec_" + name
}

func ldd(
	deps map[string]string, name string, look func(name string) (string, error),
) error {
	path, err := look(name)
	if err != nil {
		return err
	}
	deps[name] = path
	e, err := elf.Open(path)
	if err != nil {
		return err
	}
	l, err := e.ImportedLibraries()
	if err != nil {
		return err
	}
	rpaths, err := e.DynString(elf.DT_RPATH)
	if err != nil {
		return err
	}
	runpaths, err := e.DynString(elf.DT_RUNPATH)
	if err != nil {
		return err
	}
	if err := e.Close(); err != nil {
		return err
	}
	for _, lib := range l {
		if _, ok := deps[lib]; ok {
			continue
		}
		if err := ldd(deps, lib, func(name string) (string, error) {
			return lookLibrary(name, rpaths, runpaths)
		}); err != nil {
			return err
		}
	}
	return nil
}

var ldPaths = []string{"/lib", "/usr/local/lib", "/usr/lib"}

func init() {
	// TODO: os.ReadFile(fmt.Sprintf("/etc/ld-musl-$(ARCH).path"))
	switch runtime.GOARCH {
	case "amd64":
		ldPaths = append(ldPaths,
			"/lib64",
			"/usr/lib64",
			"/usr/local/lib/x86_64-linux-gnu",
			"/lib/x86_64-linux-gnu",
			"/usr/lib/x86_64-linux-gnu",
		)
	case "arm64":
		ldPaths = append(ldPaths,
			"/usr/local/lib/aarch64-linux-gnu",
			"/lib/aarch64-linux-gnu",
			"/usr/lib/aarch64-linux-gnu",
		)
	}
	if s := os.Getenv("LD_LIBRARY_PATH"); s != "" {
		ldPaths = append(ldPaths, filepath.SplitList(s)...)
	}
}

func lookLibrary(name string, rpaths, runpaths []string) (string, error) {
	for _, dir := range append(rpaths, append(ldPaths, runpaths...)...) {
		path := filepath.Join(dir, name)
		stat, err := os.Stat(path)
		if err != nil {
			if os.IsNotExist(err) {
				continue
			}
			return "", err
		}
		if !stat.IsDir() {
			return path, nil
		}
	}
	return "", fmt.Errorf("dynamic library %s not found", name)
}