// Copyright 2020 The Monogon Project Authors.
//
// SPDX-License-Identifier: Apache-2.0
//
// 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.

// gotoolwrap is a tiny wrapper used to run executables that expect a standard
// Go development environment setup to act on the monogon workspace. Notably,
// it's used in Bazel rules to run tools like gofmt when invoked by some kinds
// of code generation tools.
//
// Usage: ./gotoolwrap executable arg1 arg2
//
// gotoolwrap expects the following environment variables to be set (and unsets
// them before calling the given executable):
//  - GOTOOLWRAP_GOPATH: A synthetic GOPATH, eg. one generated by rules_go's
//  				     go_path target.
//  - GOTOOLWRAP_GOROOT: A Go SDK's GOROOT, eg. one from rules_go's GoSDK
//                       provider.
//  - GOTOOLWRAP_COPYOUT: Pairs of source:destination elements separated by ;.
//                        Files given in this list will be copied from source
//                        (relative to the GOPATH sources) to the destination
//                        (absolute path). This is primarily useful when dealing
//                        with code generators which insist on putting their
//                        generated files within the GOPATH, and allows copying
//                        of said files to a declared output file.
//
// gotoolwrap will set PATH to contain GOROOT/bin, and set GOPATH and GOROOT as
// resolved, absolute paths. Absolute paths are expected by tools like 'gofmt'.

package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

func main() {
	gopath := os.Getenv("GOTOOLWRAP_GOPATH")
	if gopath == "" {
		log.Fatal("GOTOOLWRAP_GOPATH must be set")
	}

	goroot := os.Getenv("GOTOOLWRAP_GOROOT")
	if goroot == "" {
		log.Fatal("GOTOOLWRAP_GOROOT must be set")
	}

	if len(os.Args) < 2 {
		log.Fatalf("No command specified")
	}

	// Resolve gopath and goroot to absolute paths.
	gopathAbs, err := filepath.Abs(gopath)
	if err != nil {
		log.Fatalf("Abs(%q): %v", gopath, err)
	}
	gorootAbs, err := filepath.Abs(goroot)
	if err != nil {
		log.Fatalf("Abs(%q): %v", goroot, err)
	}

	// Ensure the resolved GOROOT has a bin/go and bin/gofmt.
	gorootBin := filepath.Join(gorootAbs, "bin")
	stat, err := os.Stat(gorootBin)
	if err != nil {
		log.Fatalf("Could not stat $GOTOOLWRAP_GOROOT/bin (%q): %v", gorootBin, err)
	}
	if !stat.IsDir() {
		log.Fatalf("$GOTOOLWRAP_GOROOT/bin (%q) is not a directory", gorootBin)
	}
	// We list all files inside so that we can print them to the user for
	// debugging purposes if that's not the case.
	binFiles := make(map[string]bool)
	files, err := os.ReadDir(gorootBin)
	if err != nil {
		log.Fatalf("Could not read dir $GOTOOLWRAP_GOROOT/bin (%q): %v", gorootBin, err)
	}
	for _, f := range files {
		if f.IsDir() {
			continue
		}
		binFiles[f.Name()] = true
	}
	if !binFiles["go"] || !binFiles["gofmt"] {
		msg := "no files"
		if len(binFiles) > 0 {
			var names []string
			for name := range binFiles {
				names = append(names, fmt.Sprintf("%q", name))
			}
			msg = fmt.Sprintf(": %s", strings.Join(names, ", "))
		}
		log.Fatalf("$GOTOOLWRAP_GOROOT/bin (%q) does not contain go and/or gofmt, found %s", gorootBin, msg)
	}

	// Make new PATH.
	path := os.Getenv("PATH")
	if path == "" {
		path = gorootBin
	} else {
		path = fmt.Sprintf("%s:%s", gorootBin, path)
	}

	cmd := exec.Command(os.Args[1], os.Args[2:]...)

	// Copy current env into command's env, filtering out GOTOOLWRAP env vars
	// and PATH (which we set ourselves).
	for _, v := range os.Environ() {
		if strings.HasPrefix(v, "GOTOOLWRAP_GOROOT=") {
			continue
		}
		if strings.HasPrefix(v, "GOTOOLWRAP_GOPATH=") {
			continue
		}
		if strings.HasPrefix(v, "GOTOOLWRAP_COPYOUT=") {
			continue
		}
		if strings.HasPrefix(v, "PATH=") {
			continue
		}
		cmd.Env = append(cmd.Env, v)
	}
	// Create a temporary cache directory and remove everything afterwards.
	// Based on code from
	// @io_bazel_rules_go//go/tools/builders:stdliblist.go L174
	tempDir, err := os.MkdirTemp("", "gocache")
	if err != nil {
		log.Fatalf("Cannot create temporary directory: %v", err)
	}
	defer os.RemoveAll(tempDir)

	cmd.Env = append(cmd.Env,
		"GOROOT="+gorootAbs,
		"GOPATH="+gopathAbs,
		"PATH="+path,
		"GO111MODULE=off",
		"GOCACHE="+tempDir,
	)

	// Run the command interactively.
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	if err := cmd.Run(); err != nil {
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			os.Exit(exitErr.ExitCode())
		} else {
			log.Fatalf("Could not run %q: %v", os.Args[1], err)
		}
	}

	copyout := os.Getenv("GOTOOLWRAP_COPYOUT")
	if len(copyout) != 0 {
		for _, pair := range strings.Split(copyout, ";") {
			parts := strings.Split(pair, ":")
			if len(parts) != 2 {
				log.Fatalf("GOTOOL_COPYOUT invalid pair: %q", pair)
			}
			from := filepath.Join(gopathAbs, "src", parts[0])
			to := parts[1]
			log.Printf("gotoolwrap: Copying %s to %s...", from, to)
			data, err := os.ReadFile(from)
			if err != nil {
				log.Fatalf("gotoolwrap: read failed: %v", err)
			}
			err = os.MkdirAll(filepath.Dir(to), 0755)
			if err != nil {
				log.Fatalf("gotoolwrap: mkdir failed: %v", err)
			}
			err = os.WriteFile(to, data, 0644)
			if err != nil {
				log.Fatalf("gotoolwrap: write failed: %v", err)
			}
		}
	}
}