package core

import (
	"strings"
)

// DagResolveEnv resolves environment variable references in a DagDefinition.
// It parses environ (KEY=VALUE pairs from os.Environ()), merges with dag.Env
// (dag.Env takes precedence), and substitutes ${VAR} and $VAR in Command,
// Script, Dir, Args, and Env values of each step. Also substitutes in
// DagDefinition.WorkingDir. Returns a new DagDefinition without mutating the original.
func DagResolveEnv(dag DagDefinition, environ []string) DagDefinition {
	// Parse environ into a base map.
	base := parseEnviron(environ)

	// Merge dag.Env over base (dag.Env has precedence).
	dagEnv := mergeMaps(base, dag.Env)

	// Build new DagDefinition (shallow copy, then replace fields).
	result := dag

	// Substitute in WorkingDir.
	result.WorkingDir = substitute(dag.WorkingDir, dagEnv)

	// Substitute in Env values of the DAG itself.
	result.Env = substituteMap(dag.Env, dagEnv)

	// Resolve steps.
	resolvedSteps := make([]DagStep, len(dag.Steps))
	for i, step := range dag.Steps {
		// Merge dag env with step env (step env has precedence).
		stepEnv := mergeMaps(dagEnv, step.Env)

		rs := step
		rs.Command = substitute(step.Command, stepEnv)
		rs.Script = substitute(step.Script, stepEnv)
		rs.Dir = substitute(step.Dir, stepEnv)

		if len(step.Args) > 0 {
			args := make([]string, len(step.Args))
			for j, arg := range step.Args {
				args[j] = substitute(arg, stepEnv)
			}
			rs.Args = args
		}

		rs.Env = substituteMap(step.Env, stepEnv)
		resolvedSteps[i] = rs
	}
	result.Steps = resolvedSteps

	return result
}

// parseEnviron parses a slice of "KEY=VALUE" strings into a map.
func parseEnviron(environ []string) map[string]string {
	m := make(map[string]string, len(environ))
	for _, kv := range environ {
		idx := strings.IndexByte(kv, '=')
		if idx < 0 {
			continue
		}
		m[kv[:idx]] = kv[idx+1:]
	}
	return m
}

// mergeMaps returns a new map with all entries from base, then overlaid with overlay.
func mergeMaps(base, overlay map[string]string) map[string]string {
	result := make(map[string]string, len(base)+len(overlay))
	for k, v := range base {
		result[k] = v
	}
	for k, v := range overlay {
		result[k] = v
	}
	return result
}

// substituteMap applies substitute to all values in a map.
func substituteMap(m map[string]string, env map[string]string) map[string]string {
	if len(m) == 0 {
		return m
	}
	result := make(map[string]string, len(m))
	for k, v := range m {
		result[k] = substitute(v, env)
	}
	return result
}

// substitute replaces ${VAR} and $VAR occurrences in s using the env map.
func substitute(s string, env map[string]string) string {
	if s == "" || !strings.Contains(s, "$") {
		return s
	}

	var b strings.Builder
	i := 0
	for i < len(s) {
		if s[i] != '$' {
			b.WriteByte(s[i])
			i++
			continue
		}

		// Found '$'
		i++
		if i >= len(s) {
			b.WriteByte('$')
			break
		}

		if s[i] == '{' {
			// ${VAR} form.
			i++
			end := strings.IndexByte(s[i:], '}')
			if end < 0 {
				// No closing brace — write as-is.
				b.WriteString("${")
				continue
			}
			varName := s[i : i+end]
			i += end + 1
			if val, ok := env[varName]; ok {
				b.WriteString(val)
			} else {
				b.WriteString("${")
				b.WriteString(varName)
				b.WriteByte('}')
			}
		} else if isEnvVarChar(s[i]) {
			// $VAR form (bare variable name).
			start := i
			for i < len(s) && isEnvVarChar(s[i]) {
				i++
			}
			varName := s[start:i]
			if val, ok := env[varName]; ok {
				b.WriteString(val)
			} else {
				b.WriteByte('$')
				b.WriteString(varName)
			}
		} else {
			// Not a valid variable start — keep '$' and continue.
			b.WriteByte('$')
		}
	}

	return b.String()
}

// isEnvVarChar returns true for characters valid in environment variable names.
func isEnvVarChar(c byte) bool {
	return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '_'
}
