package profile

import (
	"bufio"
	"context"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/desktop/internal/execruntimeenv"
)

var (
	// envExportRegex matches shell export statements: export KEY=VALUE,
	// export KEY="VALUE", export KEY='VALUE'.
	envExportRegex = regexp.MustCompile(`^export\s+([A-Za-z_][A-Za-z0-9_]*)=(.*)$`)

	// shellFunctionDefinitionRegex matches common shell function definitions.
	// Static profile scanning must ignore exports inside function bodies because
	// defining a helper such as setproxy() does not execute it.
	shellFunctionDefinitionRegex = regexp.MustCompile(`^(?:[A-Za-z_][A-Za-z0-9_]*\s*\(\s*\)|function\s+[A-Za-z_][A-Za-z0-9_]*(?:\s*\(\s*\))?)\s*(?:\{|(?:#.*)?$)`)
)

const (
	shellPathStart   = "__AGENTTEAM_SHELL_PATH_START__"
	shellPathEnd     = "__AGENTTEAM_SHELL_PATH_END__"
	shellEnvStart    = "__AGENTTEAM_SHELL_ENV_START__"
	shellEnvEnd      = "__AGENTTEAM_SHELL_ENV_END__"
	shellPathTimeout = 2 * time.Second
)

var shellEnvProcessKeys = map[string]struct{}{
	"_":      {},
	"OLDPWD": {},
	"PATH":   {},
	"PWD":    {},
	"SHLVL":  {},
}

// shellProfiles returns the list of shell profile files to scan,
// ordered by preference for the current platform.
func shellProfiles(homeDir string) []string {
	return shellProfilesForGOOS(runtime.GOOS, homeDir)
}

func shellProfilesForGOOS(goos, homeDir string) []string {
	switch goos {
	case "darwin":
		return []string{
			filepath.Join(homeDir, ".zshenv"),
			filepath.Join(homeDir, ".zprofile"),
			filepath.Join(homeDir, ".zshrc"),
			filepath.Join(homeDir, ".bash_profile"),
			filepath.Join(homeDir, ".bashrc"),
			filepath.Join(homeDir, ".profile"),
		}
	case "linux":
		return []string{
			filepath.Join(homeDir, ".bashrc"),
			filepath.Join(homeDir, ".profile"),
			filepath.Join(homeDir, ".zshenv"),
			filepath.Join(homeDir, ".zshrc"),
			filepath.Join(homeDir, ".bash_profile"),
		}
	default: // windows and others
		return nil
	}
}

// LoadEnvFromShellProfiles scans the user's shell profile files for export
// statements and records runtime child-process environment fallbacks. It does
// not mutate the daemon process environment.
func LoadEnvFromShellProfiles() {
	homeDir, err := os.UserHomeDir()
	if err != nil {
		slog.Debug("shellenv: cannot resolve home directory", "error", err)
		return
	}

	profiles := shellProfiles(homeDir)
	if len(profiles) == 0 {
		return
	}

	loaded := 0
	for _, p := range profiles {
		env, err := parseProfileEnvs(p)
		if err != nil {
			if !os.IsNotExist(err) {
				slog.Debug("shellenv: error reading profile", "path", p, "error", err)
			}
			continue
		}
		loaded += runtimeconfig.AddShellEnvForRuntimeResolution(env)
	}

	if loaded > 0 {
		slog.Info("shellenv: loaded environment variables from shell profiles", "count", loaded)
	}
}

// LoadEnvFromShell asks the user's shell for its exported environment and
// records exported runtime environment variables for child processes. It does
// not mutate the daemon process environment.
func LoadEnvFromShell() {
	env, ok := resolveShellEnv(shellPathTimeout)
	if !ok {
		return
	}

	runtimeEnv := make(map[string]string, len(env))
	for key, value := range env {
		if !shouldImportShellEnvKey(key) {
			continue
		}
		runtimeEnv[key] = value
	}
	runtimeconfig.SetShellEnvForRuntimeResolution(runtimeEnv)
	if len(runtimeEnv) > 0 {
		slog.Info("shellenv: loaded runtime environment variables from user shell", "count", len(runtimeEnv))
	}
}

// LoadPathFromShell asks the user's shell to compute PATH and makes the result
// available to runtime detection and launch. It does not write user files or
// mutate the daemon process PATH.
func LoadPathFromShell() {
	shellPath, ok := resolveShellPath(shellPathTimeout)
	if !ok {
		return
	}

	separator := string(os.PathListSeparator)
	runtimeconfig.SetShellPathForRuntimeResolution(shellPath)
	slog.Info("shellenv: loaded PATH from user shell", "shell_entries", len(splitPathList(shellPath, separator)))
}

type shellPathCommand struct {
	command string
	args    []string
	parse   func(string) string
}

var shellPathCommandOutput = func(ctx context.Context, command string, args ...string) ([]byte, error) {
	cmd := exec.CommandContext(ctx, command, args...)
	cmd.Env = execruntimeenv.FilterInherited(os.Environ())
	configureShellProbeProcess(cmd)
	return cmd.CombinedOutput()
}

func resolveShellPath(timeout time.Duration) (string, bool) {
	commands := shellPathCommands()
	if len(commands) == 0 {
		return "", false
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	for _, candidate := range commands {
		if ctx.Err() != nil {
			break
		}
		out, err := shellPathCommandOutput(ctx, candidate.command, candidate.args...)
		if err != nil {
			continue
		}
		parse := candidate.parse
		if parse == nil {
			parse = extractShellPath
		}
		if path := parse(string(out)); path != "" {
			return path, true
		}
	}
	return "", false
}

func resolveShellEnv(timeout time.Duration) (map[string]string, bool) {
	return resolveShellEnvFromCommands(timeout, shellEnvCommands())
}

func resolveShellEnvFromCommands(timeout time.Duration, commands []shellPathCommand) (map[string]string, bool) {
	if len(commands) == 0 {
		return nil, false
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	for _, candidate := range commands {
		if ctx.Err() != nil {
			break
		}
		out, err := shellPathCommandOutput(ctx, candidate.command, candidate.args...)
		if err != nil {
			continue
		}
		if env := extractShellEnv(string(out)); len(env) > 0 {
			return env, true
		}
	}
	return nil, false
}

func shellPathCommands() []shellPathCommand {
	return shellPathCommandsForGOOS(runtime.GOOS, os.Getenv("SHELL"))
}

func shellEnvCommands() []shellPathCommand {
	return shellEnvCommandsForGOOS(runtime.GOOS, os.Getenv("SHELL"))
}

func shellPathCommandsForGOOS(goos, shell string) []shellPathCommand {
	if goos == "windows" {
		return windowsShellPathCommands()
	}
	return unixShellPathCommands(goos, shell)
}

func unixShellPathCommands(goos, shell string) []shellPathCommand {
	shells := make([]string, 0, 4)
	if trimmed := strings.TrimSpace(shell); trimmed != "" {
		shells = append(shells, trimmed)
	}
	switch goos {
	case "darwin":
		shells = append(shells, "/bin/zsh", "/bin/bash", "/bin/sh")
	default:
		shells = append(shells, "/bin/bash", "/bin/zsh", "/bin/sh")
	}
	shells = dedupeShells(shells)

	var commands []shellPathCommand
	for _, sh := range shells {
		script := posixShellPathProbeScript()
		if shellLooksLikeFish(sh) {
			script = fishShellPathProbeScript()
		}
		// Try interactive-login first because zsh/fnm/nvm setups commonly live
		// in .zshrc, while login-only shells read .zprofile/.bash_profile.
		for _, flag := range []string{"-ilc", "-lc", "-ic", "-c"} {
			commands = append(commands, shellPathCommand{
				command: sh,
				args:    []string{flag, script},
				parse:   extractShellPath,
			})
		}
	}
	return commands
}

func shellEnvCommandsForGOOS(goos, shell string) []shellPathCommand {
	if goos == "windows" {
		return windowsShellEnvCommands()
	}
	shells := make([]string, 0, 4)
	if trimmed := strings.TrimSpace(shell); trimmed != "" {
		shells = append(shells, trimmed)
	}
	switch goos {
	case "darwin":
		shells = append(shells, "/bin/zsh", "/bin/bash", "/bin/sh")
	default:
		shells = append(shells, "/bin/bash", "/bin/zsh", "/bin/sh")
	}
	shells = dedupeShells(shells)

	script := shellEnvProbeScript()
	var commands []shellPathCommand
	for _, sh := range shells {
		for _, flag := range []string{"-ilc", "-lc", "-ic", "-c"} {
			commands = append(commands, shellPathCommand{
				command: sh,
				args:    []string{flag, script},
				parse:   nil,
			})
		}
	}
	return commands
}

func windowsShellPathCommands() []shellPathCommand {
	powerShellScript := windowsPowerShellPathProbeScript()
	return []shellPathCommand{
		{command: "pwsh.exe", args: []string{"-NoLogo", "-Command", powerShellScript}, parse: extractShellPath},
		{command: "powershell.exe", args: []string{"-NoLogo", "-Command", powerShellScript}, parse: extractShellPath},
		{command: "cmd.exe", args: []string{"/s", "/c", "set PATH"}, parse: extractCmdPath},
	}
}

func windowsShellEnvCommands() []shellPathCommand {
	powerShellScript := windowsPowerShellEnvProbeScript()
	return []shellPathCommand{
		{command: "pwsh.exe", args: []string{"-NoLogo", "-Command", powerShellScript}},
		{command: "powershell.exe", args: []string{"-NoLogo", "-Command", powerShellScript}},
	}
}

func dedupeShells(shells []string) []string {
	result := make([]string, 0, len(shells))
	seen := make(map[string]struct{})
	for _, shell := range shells {
		shell = strings.TrimSpace(shell)
		if shell == "" {
			continue
		}
		key := shell
		if runtime.GOOS == "windows" {
			key = strings.ToLower(key)
		}
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		result = append(result, shell)
	}
	return result
}

func shellLooksLikeFish(shell string) bool {
	name := strings.ToLower(filepath.Base(strings.TrimSpace(shell)))
	return name == "fish" || name == "fish.exe"
}

func posixShellPathProbeScript() string {
	return "printf '\\n" + shellPathStart + "%s" + shellPathEnd + "\\n' \"$PATH\""
}

func fishShellPathProbeScript() string {
	return "printf '\\n" + shellPathStart + "%s" + shellPathEnd + "\\n' (string join : $PATH)"
}

func shellEnvProbeScript() string {
	return "printf '\\n" + shellEnvStart + "\\n'; env; printf '" + shellEnvEnd + "\\n'"
}

func windowsPowerShellPathProbeScript() string {
	return "[Console]::Out.Write(\"`n" + shellPathStart + "\"); [Console]::Out.Write($env:Path); [Console]::Out.Write(\"" + shellPathEnd + "`n\")"
}

func windowsPowerShellEnvProbeScript() string {
	return "[Console]::Out.Write(\"`n" + shellEnvStart + "`n\"); Get-ChildItem Env: | ForEach-Object { [Console]::Out.Write($_.Name); [Console]::Out.Write(\"=\"); [Console]::Out.WriteLine($_.Value) }; [Console]::Out.Write(\"" + shellEnvEnd + "`n\")"
}

func extractShellPath(output string) string {
	start := strings.LastIndex(output, shellPathStart)
	if start < 0 {
		return ""
	}
	rest := output[start+len(shellPathStart):]
	end := strings.Index(rest, shellPathEnd)
	if end < 0 {
		return ""
	}
	return strings.TrimSpace(rest[:end])
}

func extractCmdPath(output string) string {
	for _, line := range strings.Split(output, "\n") {
		line = strings.TrimSpace(line)
		key, value, ok := strings.Cut(line, "=")
		if !ok {
			continue
		}
		if strings.EqualFold(strings.TrimSpace(key), "path") {
			return strings.TrimSpace(value)
		}
	}
	return ""
}

func extractShellEnv(output string) map[string]string {
	start := strings.LastIndex(output, shellEnvStart)
	if start < 0 {
		return nil
	}
	rest := output[start+len(shellEnvStart):]
	end := strings.Index(rest, shellEnvEnd)
	if end < 0 {
		return nil
	}

	result := make(map[string]string)
	scanner := bufio.NewScanner(strings.NewReader(rest[:end]))
	for scanner.Scan() {
		key, value, ok := strings.Cut(scanner.Text(), "=")
		if !ok {
			continue
		}
		key = strings.TrimSpace(key)
		if key == "" || !validEnvKey(key) {
			continue
		}
		result[key] = value
	}
	if len(result) == 0 {
		return nil
	}
	return result
}

func splitPathList(value, separator string) []string {
	if value == "" {
		return nil
	}
	parts := strings.Split(value, separator)
	entries := make([]string, 0, len(parts))
	for _, part := range parts {
		entry := strings.TrimSpace(part)
		if entry != "" {
			entries = append(entries, entry)
		}
	}
	return entries
}

// parseProfileEnvs reads exported environment variables from a single profile
// file for runtime child-process environment fallback.
func parseProfileEnvs(path string) (map[string]string, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	env := make(map[string]string)
	functionDepth := 0
	pendingFunctionBody := false
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())

		// Skip empty lines and comments.
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		if functionDepth > 0 {
			functionDepth += strings.Count(line, "{") - strings.Count(line, "}")
			if functionDepth < 0 {
				functionDepth = 0
			}
			continue
		}
		if pendingFunctionBody {
			pendingFunctionBody = false
			if strings.HasPrefix(line, "{") {
				functionDepth = strings.Count(line, "{") - strings.Count(line, "}")
				if functionDepth < 0 {
					functionDepth = 0
				}
				continue
			}
		}
		if shellFunctionDefinitionRegex.MatchString(line) {
			functionDepth = strings.Count(line, "{") - strings.Count(line, "}")
			if functionDepth < 0 {
				functionDepth = 0
			}
			if functionDepth == 0 && !strings.Contains(line, "{") {
				pendingFunctionBody = true
			}
			continue
		}

		matches := envExportRegex.FindStringSubmatch(line)
		if matches == nil {
			continue
		}

		key := matches[1]
		rawValue := matches[2]

		if !shouldImportShellEnvKey(key) {
			continue
		}

		value := unquoteValue(rawValue)
		env[key] = value
	}

	return env, scanner.Err()
}

// unquoteValue strips surrounding single or double quotes from a shell value.
func unquoteValue(s string) string {
	if len(s) >= 2 {
		if (s[0] == '"' && s[len(s)-1] == '"') || (s[0] == '\'' && s[len(s)-1] == '\'') {
			if unquoted, err := strconv.Unquote(s); err == nil {
				return unquoted
			}
			return s[1 : len(s)-1]
		}
	}
	return s
}

func validEnvKey(key string) bool {
	if key == "" {
		return false
	}
	for i, r := range key {
		if i == 0 {
			if r != '_' && (r < 'A' || r > 'Z') && (r < 'a' || r > 'z') {
				return false
			}
			continue
		}
		if r != '_' && (r < 'A' || r > 'Z') && (r < 'a' || r > 'z') && (r < '0' || r > '9') {
			return false
		}
	}
	return true
}

func shouldImportShellEnvKey(key string) bool {
	key = strings.TrimSpace(key)
	if !validEnvKey(key) || execruntimeenv.IsDaemonOwnedKey(key) {
		return false
	}
	if _, skip := shellEnvProcessKeys[key]; skip {
		return false
	}
	return true
}