package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"syscall"
)

var (
	EnvVarRegex     = regexp.MustCompile(`^(.*?)=(.*?)$`)
	EnvCommentRegex = regexp.MustCompile(`^[ \t]*#.*?$`)
	ExpandEnvFiles  = map[string]string{}
)

var ( // flags
	FlagEnvFilePath       = flag.String("f", ".env", "environment file")
	FlagExpandRecursively = flag.Bool("r", false, "expand references in reference definitions")

	FlagExpandEnvFiles = SliceFlag{}
)

func main() {
	flag.Var(&FlagExpandEnvFiles, "e", "files in which to expand environment variables")
	flag.Parse()

	envFile, err := os.Open(*FlagEnvFilePath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "could not open file '%s': %s\n", *FlagEnvFilePath, err.(*os.PathError).Err)
		return
	}
	defer envFile.Close()

	parseEnvFile(envFile)
	args := flag.Args()

	if len(args) == 0 {
		for _, envVar := range os.Environ() {
			fmt.Println(envVar)
		}
		return
	}

	if err := expandEnvFiles(); err != nil {
		fmt.Fprintf(os.Stderr, "could not expand files: %s\n", err.Error())
		return
	}
	defer restoreFiles()

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() { <-sigs; exit(1) }()

	cmd := exec.Command(args[0], args[1:]...)
	cmd.Env = os.Environ()

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			exit(exitErr.ExitCode())
			return
		} else {
			panic(err)
		}
	}
}

func parseEnvFile(r io.Reader) {
	s := bufio.NewScanner(r)
	l := 0

	for s.Scan() {
		l++

		matches := EnvVarRegex.FindStringSubmatch(s.Text())
		if len(matches) != 0 {
			key, value := matches[1], matches[2]

			if *FlagExpandRecursively {
				value = os.ExpandEnv(value)
			}

			os.Setenv(key, value)
		} else if s.Text() != "" && !EnvCommentRegex.MatchString(s.Text()) {
			panic(fmt.Sprintf("invalid env syntax on line %d", l))
		}
	}
}

func exit(errorCode int) {
	restoreFiles()
	os.Exit(errorCode)
}

func restoreFiles() {
	if err := restoreEnvFiles(); err != nil {
		fmt.Fprintf(os.Stderr, "could not restore expanded files: %s\n", err.Error())
	}
}