package main

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"syscall"
)

func main() {
	if len(os.Args) < 3 {
		fmt.Fprintf(os.Stderr, "usage: envdir dir child\n")
		os.Exit(111)
	}

	dir := os.Args[1]
	child := os.Args[2]
	childArgs := os.Args[2:]

	env, err := buildEnv(dir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "envdir: %s\n", err)
		os.Exit(111)
	}

	path, err := findExecutable(child)
	if err != nil {
		fmt.Fprintf(os.Stderr, "envdir: %s\n", err)
		os.Exit(111)
	}

	err = syscall.Exec(path, childArgs, env)
	fmt.Fprintf(os.Stderr, "envdir: exec: %s\n", err)
	os.Exit(111)
}

func buildEnv(dir string) ([]string, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, fmt.Errorf("unable to read directory %s: %w", dir, err)
	}

	// Start with the current environment.
	envMap := make(map[string]string)
	unset := make(map[string]bool)
	for _, e := range os.Environ() {
		k, v, _ := strings.Cut(e, "=")
		envMap[k] = v
	}

	for _, entry := range entries {
		name := entry.Name()

		// Skip subdirectories.
		if entry.IsDir() {
			continue
		}

		// Skip names containing '='.
		if strings.Contains(name, "=") {
			continue
		}

		path := filepath.Join(dir, name)
		data, err := os.ReadFile(path)
		if err != nil {
			return nil, fmt.Errorf("unable to read %s: %w", path, err)
		}

		// If the file is empty, unset the variable.
		if len(data) == 0 {
			delete(envMap, name)
			unset[name] = true
			continue
		}

		// Use only the first line.
		if i := bytes.IndexByte(data, '\n'); i >= 0 {
			data = data[:i]
		}

		// Replace NUL bytes with newlines.
		data = bytes.ReplaceAll(data, []byte{0}, []byte{'\n'})

		// Trim trailing spaces and tabs.
		value := strings.TrimRight(string(data), " \t")

		envMap[name] = value
		delete(unset, name)
	}

	env := make([]string, 0, len(envMap))
	for k, v := range envMap {
		env = append(env, k+"="+v)
	}
	return env, nil
}

func findExecutable(name string) (string, error) {
	// If name contains a slash, use it directly.
	if strings.Contains(name, "/") {
		return name, nil
	}

	// Search PATH.
	pathEnv := os.Getenv("PATH")
	for _, dir := range filepath.SplitList(pathEnv) {
		path := filepath.Join(dir, name)
		info, err := os.Stat(path)
		if err != nil {
			continue
		}
		if info.Mode().IsRegular() && info.Mode()&0111 != 0 {
			return path, nil
		}
	}
	return "", fmt.Errorf("command not found: %s", name)
}