package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"
)

func exit(err error) {
	fmt.Fprintf(os.Stderr, "[offsets] error: %s\n", err.Error())
	os.Exit(1)
}

func genBuildScript(targetOS, targetArch, goBinary, workDir string) ([]byte, error) {
	// Write a dummy program in workDir so "go build" does not complain
	dummyGoProgram := []byte("package main\n func main(){}")
	err := ioutil.WriteFile(fmt.Sprintf("%s/main.go", workDir), dummyGoProgram, os.ModePerm)
	if err != nil {
		return nil, err
	}

	// Run "go build -a -n" in workDir and capture the output. The -a flag
	// ensures that the generated build script includes steps to always
	// rebuild the runtime packages.
	cmd := exec.Command(goBinary, "build", "-a", "-n")
	cmd.Dir = workDir
	cmd.Env = append(cmd.Env, fmt.Sprintf("GOROOT=%s", os.Getenv("GOROOT")))
	cmd.Env = append(cmd.Env, fmt.Sprintf("GOOS=%s", targetOS))
	cmd.Env = append(cmd.Env, fmt.Sprintf("GOARCH=%s", targetArch))
	out, err := cmd.CombinedOutput()
	if err != nil {
		return nil, fmt.Errorf("failed to generate build script\nMore info:\n%s", out)
	}

	return out, nil
}

func patchBuildScript(script []byte, workDir, targetOS, targetArch, goBinary string) ([]byte, error) {
	// Replace $WORK with the workDir location. This is required for executing
	// build scripts generated by go 1.10
	lines := strings.Split(
		strings.Replace(string(script), "$WORK", workDir, -1),
		"\n",
	)

	// Inject os/arch and workdir to the top of the build file
	header := []string{
		fmt.Sprintf("export GOROOT=%s", os.Getenv("GOROOT")),
		fmt.Sprintf("export GOOS=%s", targetOS),
		fmt.Sprintf("export GOARCH=%s", targetArch),
		fmt.Sprintf("alias pack='%s tool pack'", goBinary),
	}
	lines = append(header, lines...)

	// We are only interested in building the runtime as this block generates
	// the asm headers. Scan the lines till we find "# runtime" comment and
	// stop at next comment
	var stopOnNextComment bool
	for lineIndex := 0; lineIndex < len(lines); lineIndex++ {
		// Ignore empty comments
		if strings.TrimSpace(lines[lineIndex]) == "#" || strings.Contains(lines[lineIndex], "# import") {
			continue
		}

		if stopOnNextComment && strings.HasPrefix(lines[lineIndex], "#") {
			return []byte(strings.Join(lines[:lineIndex], "\n")), nil
		}

		if lines[lineIndex] == "# runtime" {
			stopOnNextComment = true
		}
	}

	return nil, errors.New("generated build file does not specify -asmhdr when building the runtime")
}

func execBuildScript(script []byte, workDir string) error {
	f, err := os.Create(fmt.Sprintf("%s/build.sh", workDir))
	if err != nil {
		return err
	}

	_, err = f.Write(script)
	if err != nil {
		f.Close()
		return err
	}
	f.Close()

	cmd := exec.Command("sh", f.Name())
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to execute build script\nMore info:\n%s", out)
	}

	return nil
}

func genAsmIncludes(workDir string) ([]byte, error) {
	// Find all generated go_asm.h files and concat their conentents
	var (
		allHeaders, headers []byte
	)

	if err := filepath.Walk(workDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		if filepath.Base(path) != "go_asm.h" {
			return nil
		}

		if headers, err = ioutil.ReadFile(path); err != nil {
			return err
		}

		allHeaders = append(allHeaders, '\n')
		allHeaders = append(allHeaders, headers...)
		return nil
	}); err != nil {
		return nil, err
	}

	var includes []string
	includes = append(includes, "; vim: set ft=nasm :\n")
	includes = append(includes, fmt.Sprintf("; generated by tools/offsets at %v\n", time.Now()))

	for _, line := range strings.Split(string(allHeaders), "\n") {
		line = strings.TrimPrefix(line, "#define ")

		// We are only interested in the offsets for the g, m and stack structures
		if strings.HasPrefix(line, "g_") || strings.HasPrefix(line, "m_") || strings.HasPrefix(line, "stack_") {
			tokens := strings.Fields(line)
			if len(tokens) != 2 {
				continue
			}

			offset, err := strconv.ParseInt(tokens[1], 10, 32)
			if err != nil {
				continue
			}

			includes = append(includes,
				fmt.Sprintf("GO_%s equ 0x%x ; %d",
					strings.ToUpper(tokens[0]),
					offset,
					offset,
				),
			)
		}
	}

	// In go 1.7.x, the page size is given by the _PageSize constant whereas in
	// newer go versions it is specified at runtime and needs to be manually set
	// by our asm bootstrap code.
	if strings.Contains(runtime.Version(), "go1.7") {
		includes = append(includes,
			"; go 1.7 runtime uses a fixed 4k page size for our target arch so our bootstrap code does not need to do any extra work to set it up",
			"%define SKIP_PAGESIZE_SETUP 1",
		)
	}

	return []byte(strings.Join(includes, "\n")), nil
}

func runTool() error {
	targetOS := flag.String("target-os", "", "a valid GOOS value for generating the asm offsets")
	targetArch := flag.String("target-arch", "", "a valid GOARCH value for generating the asm offsets")
	goBinary := flag.String("go-binary", "go", "the Go binary to use")
	asmOutput := flag.String("out", "-", "a file to write the asm headers or - to output to STDOUT")
	flag.Parse()

	switch {
	case *targetOS == "":
		exit(errors.New("-target-os parameter missing"))
	case *targetArch == "":
		exit(errors.New("-target-arch parameter missing"))
	}

	workDir, err := ioutil.TempDir("", "offsets-tool")
	if err != nil {
		return err
	}
	defer os.RemoveAll(workDir)

	buildScript, err := genBuildScript(*targetOS, *targetArch, *goBinary, workDir)
	if err != nil {
		return err
	}

	buildScript, err = patchBuildScript(buildScript, workDir, *targetOS, *targetArch, *goBinary)
	if err != nil {
		return err
	}

	if err = execBuildScript(buildScript, workDir); err != nil {
		return err
	}

	asmIncludes, err := genAsmIncludes(workDir)
	if err != nil {
		return err
	}

	switch *asmOutput {
	case "-":
		fmt.Printf("%s\n", string(asmIncludes))
	default:
		if err = ioutil.WriteFile(*asmOutput, asmIncludes, os.ModePerm); err != nil {
			return err
		}
	}

	return nil
}

func main() {
	if err := runTool(); err != nil {
		exit(err)
	}
}