file.go

package env

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"regexp"
	"sort"
	"strings"
)

const doubleQuoteSpecialChars = "\\\n\r\"!$`"

// Load reads environment variables from .env files and sets them in the current process.
// If no filenames are provided, it defaults to loading ".env" from the current directory.
//
// Multiple files can be loaded in order:
//
//	err := env.Load("base.env", "local.env")
//
// Important: Load will NOT override environment variables that are already set.
// This allows .env files to provide defaults while respecting existing environment configuration.
// Use Overload if you need to override existing variables.
//
// The function returns an error if any file cannot be read or parsed.
// Call this early in your program, typically in main().
func Load(filenames ...string) (err error) {
	filenames = filenamesOrDefault(filenames)

	for _, filename := range filenames {
		err = loadFile(filename, false)
		if err != nil {
			return // return early on a spazout
		}
	}
	return
}

// MustLoad reads environment variables from .env files and sets them in the current process.
// If no filenames are provided, it defaults to loading ".env" from the current directory.
//
// This function behaves exactly like Load, except it panics on any error instead of returning it.
// Use this when .env file loading is critical for your application to function.
//
// Multiple files can be loaded in order:
//
//	env.MustLoad("base.env", "local.env")
//
// Like Load, this will NOT override environment variables that are already set.
func MustLoad(filenames ...string) {
	filenames = filenamesOrDefault(filenames)

	for _, filename := range filenames {
		err := loadFile(filename, false)
		if err != nil {
			panic(err)
		}
	}
}

// Overload reads environment variables from .env files and sets them in the current process,
// overriding any existing environment variables.
// If no filenames are provided, it defaults to loading ".env" from the current directory.
//
// Unlike Load, this function WILL override environment variables that are already set.
// Use this when you want .env files to take precedence over existing environment configuration.
//
// Multiple files can be loaded in order:
//
//	err := env.Overload("base.env", "production.env")
//
// The function returns an error if any file cannot be read or parsed.
func Overload(filenames ...string) (err error) {
	filenames = filenamesOrDefault(filenames)

	for _, filename := range filenames {
		err = loadFile(filename, true)
		if err != nil {
			return // return early on a spazout
		}
	}
	return
}

// MustOverload reads environment variables from .env files and sets them in the current process,
// overriding any existing environment variables.
// If no filenames are provided, it defaults to loading ".env" from the current directory.
//
// This function behaves exactly like Overload, except it panics on any error instead of returning it.
// Use this when .env file loading is critical and should override existing configuration.
//
// Multiple files can be loaded in order:
//
//	env.MustOverload("base.env", "production.env")
func MustOverload(filenames ...string) {
	filenames = filenamesOrDefault(filenames)

	for _, filename := range filenames {
		err := loadFile(filename, true)
		if err != nil {
			panic(err)
		}
	}
}

// Read parses .env files and returns the key-value pairs as a map,
// without setting them as environment variables in the current process.
//
// This is useful when you want to inspect or manipulate the values before
// applying them, or when you want to use the values in a different way.
// If no filenames are provided, it defaults to reading ".env" from the current directory.
//
// Multiple files can be read, with later files overriding earlier ones:
//
//	envMap, err := env.Read("base.env", "local.env")
//	if err != nil {
//		log.Fatal(err)
//	}
//	fmt.Printf("DATABASE_URL=%s\n", envMap["DATABASE_URL"])
func Read(filenames ...string) (envMap map[string]string, err error) {
	filenames = filenamesOrDefault(filenames)
	envMap = make(map[string]string)

	for _, filename := range filenames {
		individualEnvMap, individualErr := readFile(filename)

		if individualErr != nil {
			err = individualErr
			return // return early on a spazout
		}

		for key, value := range individualEnvMap {
			envMap[key] = value
		}
	}

	return
}

// ParseIO reads environment variables from an io.Reader in .env format.
// It parses the content and returns a map of key-value pairs.
//
// This function supports the standard .env file format including:
//   - Comments (lines starting with #)
//   - Quoted values (both single and double quotes)
//   - Variable expansion (${VAR} and $VAR syntax)
//   - Multiline values
//   - Both KEY=value and KEY: value formats
//
// It does not set any environment variables; it only parses and returns the data.
func ParseIO(r io.Reader) (envMap map[string]string, err error) {
	envMap = make(map[string]string)

	var lines []string
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	if err = scanner.Err(); err != nil {
		return
	}

	for _, fullLine := range lines {
		if !isIgnoredLine(fullLine) {
			var key, value string
			key, value, err = parseLine(fullLine, envMap)

			if err != nil {
				return
			}
			envMap[key] = value
		}
	}
	return
}

// Unmarshal parses environment variables from a string in .env format.
// It returns a map of key-value pairs without setting any environment variables.
//
// This is a convenience function that wraps ParseIO with a strings.NewReader.
// See ParseIO for details on the supported .env file format.
func Unmarshal(str string) (envMap map[string]string, err error) {
	return ParseIO(strings.NewReader(str))
}

// Write serializes a map of environment variables to a .env file.
// The output file will contain KEY="VALUE" pairs, one per line,
// with keys sorted alphabetically and values properly escaped.
//
// This function will create or overwrite the specified file.
func Write(envMap map[string]string, filename string) error {
	content, error := Marshal(envMap)
	if error != nil {
		return error
	}
	file, error := os.Create(filename)
	if error != nil {
		return error
	}
	_, err := file.WriteString(content)
	return err
}

// Marshal converts a map of environment variables to .env file format.
// It returns a string with KEY="VALUE" pairs, one per line,
// with keys sorted alphabetically and values properly escaped with backslashes.
//
// This function does not write to any file; use Write to save the output to disk.
func Marshal(envMap map[string]string) (string, error) {
	lines := make([]string, 0, len(envMap))
	for k, v := range envMap {
		lines = append(lines, fmt.Sprintf(`%s="%s"`, k, doubleQuoteEscape(v)))
	}
	sort.Strings(lines)
	return strings.Join(lines, "\n"), nil
}

func filenamesOrDefault(filenames []string) []string {
	if len(filenames) == 0 {
		return []string{".env"}
	}
	return filenames
}

func loadFile(filename string, overload bool) error {
	envMap, err := readFile(filename)
	if err != nil {
		return err
	}

	for key, value := range envMap {
		if _, exists := os.LookupEnv(key); !exists || overload {
			os.Setenv(key, value)
		}
	}

	return nil
}

func readFile(filename string) (envMap map[string]string, err error) {
	file, err := os.Open(filename)
	if err != nil {
		return
	}
	defer file.Close()

	return ParseIO(file)
}

func parseLine(line string, envMap map[string]string) (key string, value string, err error) {
	if len(line) == 0 {
		err = errors.New("zero length string")
		return
	}

	// ditch the comments (but keep quoted hashes)
	if strings.Contains(line, "#") {
		segmentsBetweenHashes := strings.Split(line, "#")
		quotesAreOpen := false
		var segmentsToKeep []string
		for _, segment := range segmentsBetweenHashes {
			if strings.Count(segment, "\"") == 1 || strings.Count(segment, "'") == 1 {
				if quotesAreOpen {
					quotesAreOpen = false
					segmentsToKeep = append(segmentsToKeep, segment)
				} else {
					quotesAreOpen = true
				}
			}

			if len(segmentsToKeep) == 0 || quotesAreOpen {
				segmentsToKeep = append(segmentsToKeep, segment)
			}
		}

		line = strings.Join(segmentsToKeep, "#")
	}

	firstEquals := strings.Index(line, "=")
	firstColon := strings.Index(line, ":")
	splitString := strings.SplitN(line, "=", 2)
	if firstColon != -1 && (firstColon < firstEquals || firstEquals == -1) {
		//this is a yaml-style line
		splitString = strings.SplitN(line, ":", 2)
	}

	if len(splitString) != 2 {
		err = errors.New("can't separate key from value")
		return
	}

	// Parse the key
	re := regexp.MustCompile(`^\s*(?:export\s+)?(.*?)\s*$`)
	key = re.ReplaceAllString(splitString[0], "$1")

	// Parse the value
	value = parseValue(splitString[1], envMap)
	return
}

func parseValue(value string, envMap map[string]string) string {

	// trim
	value = strings.Trim(value, " ")

	// check if we've got quoted values or possible escapes
	if len(value) > 1 {
		rs := regexp.MustCompile(`\A'(.*)'\z`)
		singleQuotes := rs.FindStringSubmatch(value)

		rd := regexp.MustCompile(`\A"(.*)"\z`)
		doubleQuotes := rd.FindStringSubmatch(value)

		if singleQuotes != nil || doubleQuotes != nil {
			// pull the quotes off the edges
			value = value[1 : len(value)-1]
		}

		if doubleQuotes != nil {
			// expand newlines
			escapeRegex := regexp.MustCompile(`\\.`)
			value = escapeRegex.ReplaceAllStringFunc(value, func(match string) string {
				c := strings.TrimPrefix(match, `\`)
				switch c {
				case "n":
					return "\n"
				case "r":
					return "\r"
				default:
					return match
				}
			})
			// unescape characters
			e := regexp.MustCompile(`\\([^$])`)
			value = e.ReplaceAllString(value, "$1")
		}

		if singleQuotes == nil {
			value = expandVariables(value, envMap)
		}
	}

	return value
}

func expandVariables(v string, m map[string]string) string {
	r := regexp.MustCompile(`(\\)?(\$)(\()?\{?([A-Z0-9_]+)?\}?`)

	return r.ReplaceAllStringFunc(v, func(s string) string {
		submatch := r.FindStringSubmatch(s)

		if submatch == nil {
			return s
		}
		if submatch[1] == "\\" || submatch[2] == "(" {
			return submatch[0][1:]
		} else if submatch[4] != "" {
			return m[submatch[4]]
		}
		return s
	})
}

func isIgnoredLine(line string) bool {
	trimmedLine := strings.TrimSpace(line)
	return len(trimmedLine) == 0 || strings.HasPrefix(trimmedLine, "#")
}

func doubleQuoteEscape(line string) string {
	for _, c := range doubleQuoteSpecialChars {
		toReplace := "\\" + string(c)
		if c == '\n' {
			toReplace = `\n`
		}
		if c == '\r' {
			toReplace = `\r`
		}
		line = strings.ReplaceAll(line, string(c), toReplace)
	}
	return line
}