package profile

import (
	"context"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

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

func TestUnquoteValue(t *testing.T) {
	tests := []struct {
		input string
		want  string
	}{
		{`"hello"`, "hello"},
		{`'hello'`, "hello"},
		{`hello`, "hello"},
		{`"hello`, `"hello`},
		{`hello"`, `hello"`},
		{`""`, ""},
		{`''`, ""},
		{`"a'b"`, "a'b"},
		{`'a"b'`, `a"b`},
	}

	for _, tc := range tests {
		got := unquoteValue(tc.input)
		if got != tc.want {
			t.Errorf("unquoteValue(%q) = %q, want %q", tc.input, got, tc.want)
		}
	}
}

func TestShellProbeEnvironmentStripsInheritedExecCapability(t *testing.T) {
	const helperKey = "AGENTTEAM_TEST_SHELL_PROBE_HELPER"
	t.Setenv("AGENTTEAM_EXEC_BASE_URL", "http://127.0.0.1:1")
	t.Setenv("AGENTTEAM_EXEC_TOKEN", "ambient-bearer")
	t.Setenv("AGENTTEAM_EXEC_TOKEN_FILE", "/private/token")
	t.Setenv("AGENTTEAM_SHELL_PROBE_SAFE", "kept")
	t.Setenv(helperKey, "1")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	output, err := shellPathCommandOutput(ctx, os.Args[0], "-test.run=^TestShellProbeEnvironmentHelper$")
	if err != nil {
		t.Fatalf("shell probe helper inherited reserved environment: %v: %s", err, output)
	}
}

func TestShellProbeEnvironmentHelper(t *testing.T) {
	if os.Getenv("AGENTTEAM_TEST_SHELL_PROBE_HELPER") != "1" {
		return
	}
	for _, key := range []string{"AGENTTEAM_EXEC_BASE_URL", "AGENTTEAM_EXEC_TOKEN", "AGENTTEAM_EXEC_TOKEN_FILE"} {
		if _, ok := os.LookupEnv(key); ok {
			os.Exit(2)
		}
	}
	if os.Getenv("AGENTTEAM_SHELL_PROBE_SAFE") != "kept" {
		os.Exit(3)
	}
	os.Exit(0)
}

func TestShouldImportShellEnvKey(t *testing.T) {
	tests := []struct {
		key  string
		want bool
	}{
		{"CUSTOM_LLM_REGION", true},
		{"AZURE_OPENAI_API_VERSION", true},
		{"RUNTIME_PROFILE", true},
		{"OPENAI_API_KEY", true},
		{"AGENTTEAM_EXEC_TOKEN", false},
		{"agentteam_exec_token_file", false},
		{"PATH", false},
		{"PWD", false},
		{"SHLVL", false},
		{"BAD-KEY", false},
		{"1INVALID", false},
	}

	for _, tc := range tests {
		got := shouldImportShellEnvKey(tc.key)
		if got != tc.want {
			t.Errorf("shouldImportShellEnvKey(%q) = %v, want %v", tc.key, got, tc.want)
		}
	}
}

func TestParseProfileEnvs(t *testing.T) {
	dir := t.TempDir()
	profilePath := filepath.Join(dir, ".testrc")

	content := strings.Join([]string{
		`# This is a comment`,
		`export LLMPROXY_OAI_KEY=test-key-123`,
		`export CUSTOM_LLM_REGION="us-east-1"`,
		`export PATH=/usr/bin:/bin`,
		`export RANDOM_VAR=value`,
		`  export RUNTIME_PROFILE=prod  `,
		``,
		`# Another comment`,
		`export AZURE_OPENAI_API_VERSION=2025-04-01-preview`,
		`alias ll='ls -la'`, // not an export
	}, "\n")

	if err := os.WriteFile(profilePath, []byte(content), 0644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}

	env, err := parseProfileEnvs(profilePath)
	if err != nil {
		t.Fatalf("parseProfileEnvs: %v", err)
	}

	if len(env) != 5 {
		t.Errorf("len(env) = %d, want 5: %#v", len(env), env)
	}

	// Verify values.
	tests := []struct {
		key, want string
	}{
		{"LLMPROXY_OAI_KEY", "test-key-123"},
		{"CUSTOM_LLM_REGION", "us-east-1"},
		{"RANDOM_VAR", "value"},
		{"RUNTIME_PROFILE", "prod"},
		{"AZURE_OPENAI_API_VERSION", "2025-04-01-preview"},
	}

	for _, tc := range tests {
		got := env[tc.key]
		if got != tc.want {
			t.Errorf("env[%q] = %q, want %q", tc.key, got, tc.want)
		}
	}
	if _, ok := env["PATH"]; ok {
		t.Fatalf("PATH should not be imported: %#v", env)
	}
}

func TestParseProfileEnvsIgnoresFunctionBodyExports(t *testing.T) {
	tests := []struct {
		name  string
		lines []string
	}{
		{
			name: "inline braces",
			lines: []string{
				`export TOP_LEVEL_VALUE=kept`,
				`setproxy() {`,
				`export HTTP_PROXY="http://127.0.0.1:8088"`,
				`export HTTPS_PROXY="http://127.0.0.1:8088"`,
				`}`,
				`function corplink_proxy {`,
				`export http_proxy="http://127.0.0.1:8888"`,
				`export ALL_PROXY="socks5://127.0.0.1:1088"`,
				`}`,
				`#setproxy`,
				`export AFTER_FUNCTION=also-kept`,
			},
		},
		{
			name: "multiline brace",
			lines: []string{
				`setproxy()`,
				`{`,
				`export HTTP_PROXY="http://127.0.0.1:8088"`,
				`}`,
				`export TOP_LEVEL_VALUE=kept`,
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			profilePath := filepath.Join(t.TempDir(), ".testrc")
			if err := os.WriteFile(profilePath, []byte(strings.Join(tc.lines, "\n")), 0644); err != nil {
				t.Fatalf("WriteFile: %v", err)
			}

			env, err := parseProfileEnvs(profilePath)
			if err != nil {
				t.Fatalf("parseProfileEnvs: %v", err)
			}
			if got := env["TOP_LEVEL_VALUE"]; got != "kept" {
				t.Fatalf("env[TOP_LEVEL_VALUE] = %q, want kept", got)
			}
			for _, key := range []string{"HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "ALL_PROXY"} {
				if _, ok := env[key]; ok {
					t.Fatalf("%s should not be imported from a shell function body: %#v", key, env)
				}
			}
		})
	}
}

func TestLoadPathFromShellCachesShellPathWithoutMutatingProcessPath(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("shell stub uses a POSIX shebang")
	}
	runtimeconfig.SetShellPathForRuntimeResolution("")
	t.Cleanup(func() {
		runtimeconfig.SetShellPathForRuntimeResolution("")
	})

	dir := t.TempDir()
	shellPath := filepath.Join(dir, "shell")
	sep := string(os.PathListSeparator)
	shellOnly := filepath.Join(dir, "fnm", "bin")
	currentOnly := filepath.Join(dir, "agentteam", "bin")
	duplicate := filepath.Join(dir, "usr", "bin")
	shellOutputPath := shellOnly + sep + duplicate
	currentPath := currentOnly + sep + duplicate
	stubShellProbeOutput(t, shellPath, "profile noise before sentinel\n"+shellPathStart+shellOutputPath+shellPathEnd+"\n")

	t.Setenv("SHELL", shellPath)
	t.Setenv("PATH", currentPath)

	LoadPathFromShell()

	if got := os.Getenv("PATH"); got != currentPath {
		t.Fatalf("PATH = %q, want unchanged %q", got, currentPath)
	}
	if got := runtimeconfig.ShellPathForRuntimeResolution(); got != shellOutputPath {
		t.Fatalf("ShellPathForRuntimeResolution() = %q, want %q", got, shellOutputPath)
	}
}

func TestLoadEnvFromShellCachesRuntimeEnvWithoutMutatingDaemonEnv(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("shell stub uses a POSIX shebang")
	}
	runtimeconfig.SetShellEnvForRuntimeResolution(nil)
	t.Cleanup(func() {
		runtimeconfig.SetShellEnvForRuntimeResolution(nil)
	})

	dir := t.TempDir()
	shellPath := filepath.Join(dir, "fish")
	processPath := filepath.Join(dir, "agentteam", "bin")
	stubShellProbeOutput(t, shellPath, strings.Join([]string{
		"profile noise before sentinel",
		shellEnvStart,
		"OPENAI_API_KEY=sk-from-fish",
		"CUSTOM_LLM_REGION=us-east-1",
		"AZURE_OPENAI_API_VERSION=2025-04-01-preview",
		"PATH=/should/not/import",
		"PWD=/should/not/import",
		shellEnvEnd,
	}, "\n"))

	for _, key := range []string{"OPENAI_API_KEY", "CUSTOM_LLM_REGION", "AZURE_OPENAI_API_VERSION"} {
		t.Setenv(key, "")
		os.Unsetenv(key)
	}
	t.Setenv("SHELL", shellPath)
	t.Setenv("PATH", processPath)

	LoadEnvFromShell()

	if got := os.Getenv("OPENAI_API_KEY"); got != "" {
		t.Fatalf("daemon OPENAI_API_KEY = %q, want empty", got)
	}
	env := runtimeconfig.ShellEnvForRuntimeResolution()
	if got := env["OPENAI_API_KEY"]; got != "sk-from-fish" {
		t.Fatalf("runtime OPENAI_API_KEY = %q, want sk-from-fish", got)
	}
	if got := env["CUSTOM_LLM_REGION"]; got != "us-east-1" {
		t.Fatalf("runtime CUSTOM_LLM_REGION = %q, want us-east-1", got)
	}
	if got := env["AZURE_OPENAI_API_VERSION"]; got != "2025-04-01-preview" {
		t.Fatalf("runtime AZURE_OPENAI_API_VERSION = %q, want 2025-04-01-preview", got)
	}
	if _, ok := env["PATH"]; ok {
		t.Fatalf("runtime env should not import PATH: %#v", env)
	}
	if _, ok := env["PWD"]; ok {
		t.Fatalf("runtime env should not import PWD: %#v", env)
	}
	if got := os.Getenv("PATH"); got != processPath {
		t.Fatalf("PATH = %q, want unchanged %q", got, processPath)
	}
}

func TestLoadPathFromShellLetsRuntimeProbeFindCodexInShellOnlyFNMPath(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("shell stub uses a POSIX shebang")
	}
	runtimeconfig.SetShellPathForRuntimeResolution("")
	oldProbe := runtimeconfig.CodexAppServerAvailable
	t.Cleanup(func() {
		runtimeconfig.SetShellPathForRuntimeResolution("")
		runtimeconfig.CodexAppServerAvailable = oldProbe
	})
	runtimeconfig.CodexAppServerAvailable = func(string) bool { return true }

	dir := t.TempDir()
	fnmBin := filepath.Join(dir, ".local", "state", "fnm_multishells", "22448_1779160692450", "bin")
	if err := os.MkdirAll(fnmBin, 0o755); err != nil {
		t.Fatalf("MkdirAll fnm bin: %v", err)
	}
	codexPath := filepath.Join(fnmBin, "codex")
	if err := os.WriteFile(codexPath, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil {
		t.Fatalf("WriteFile codex: %v", err)
	}

	shellPath := filepath.Join(dir, "shell")
	stubShellProbeOutput(t, shellPath, shellPathStart+fnmBin+shellPathEnd+"\n")

	processPath := filepath.Join(dir, "agentteam", "bin")
	if err := os.MkdirAll(processPath, 0o755); err != nil {
		t.Fatalf("MkdirAll process path: %v", err)
	}
	t.Setenv("SHELL", shellPath)
	t.Setenv("PATH", processPath)

	LoadPathFromShell()
	probe := runtimeconfig.ProbeRuntimeBinary(runtimeconfig.TypeCodex, "codex")

	if !probe.Capable {
		t.Fatalf("ProbeRuntimeBinary() capable = false, err = %v", probe.Err)
	}
	if probe.Path != codexPath {
		t.Fatalf("ProbeRuntimeBinary() path = %q, want %q", probe.Path, codexPath)
	}
	if got := os.Getenv("PATH"); got != processPath {
		t.Fatalf("PATH = %q, want unchanged %q", got, processPath)
	}
}

func stubShellProbeOutput(t *testing.T, shellPath, output string) {
	t.Helper()
	original := shellPathCommandOutput
	t.Cleanup(func() { shellPathCommandOutput = original })
	shellPathCommandOutput = func(_ context.Context, command string, args ...string) ([]byte, error) {
		if command != shellPath {
			t.Fatalf("shellPathCommandOutput command = %q, want %q", command, shellPath)
		}
		if len(args) == 0 || args[0] != "-ilc" {
			t.Fatalf("shellPathCommandOutput args = %#v, want interactive login shell", args)
		}
		return []byte(output), nil
	}
}

func TestExtractShellEnvIgnoresNoiseAndInvalidKeys(t *testing.T) {
	output := strings.Join([]string{
		"before",
		shellEnvStart,
		"OPENAI_API_KEY=sk-test",
		"OPENAI_BASE_URL=https://llm.example.com/v1",
		"BAD-KEY=value",
		"NO_EQUALS",
		"VALUE_WITH_EQUALS=a=b=c",
		shellEnvEnd,
		"after",
	}, "\n")

	got := extractShellEnv(output)
	if got["OPENAI_API_KEY"] != "sk-test" {
		t.Fatalf("OPENAI_API_KEY = %q, want sk-test", got["OPENAI_API_KEY"])
	}
	if got["OPENAI_BASE_URL"] != "https://llm.example.com/v1" {
		t.Fatalf("OPENAI_BASE_URL = %q, want https://llm.example.com/v1", got["OPENAI_BASE_URL"])
	}
	if _, ok := got["BAD-KEY"]; ok {
		t.Fatalf("invalid env key was imported: %#v", got)
	}
	if got["VALUE_WITH_EQUALS"] != "a=b=c" {
		t.Fatalf("VALUE_WITH_EQUALS = %q, want a=b=c", got["VALUE_WITH_EQUALS"])
	}
}

func TestExtractShellPathIgnoresNoise(t *testing.T) {
	output := "before\n" + shellPathStart + "/a:/b" + shellPathEnd + "\nafter\n"
	if got := extractShellPath(output); got != "/a:/b" {
		t.Fatalf("extractShellPath() = %q", got)
	}
}

func TestExtractCmdPath(t *testing.T) {
	output := "noise\r\nPath=C:\\Tools;C:\\Windows\\System32\r\nPATHEXT=.COM;.EXE\r\n"
	if got := extractCmdPath(output); got != `C:\Tools;C:\Windows\System32` {
		t.Fatalf("extractCmdPath() = %q", got)
	}
}

func TestResolveShellEnvFromWindowsPowerShellCommand(t *testing.T) {
	oldOutput := shellPathCommandOutput
	t.Cleanup(func() {
		shellPathCommandOutput = oldOutput
	})
	shellPathCommandOutput = func(_ context.Context, command string, args ...string) ([]byte, error) {
		if command != "pwsh.exe" {
			t.Fatalf("shellPathCommandOutput command = %q, want pwsh.exe", command)
		}
		if len(args) == 0 || !strings.Contains(args[len(args)-1], "Get-ChildItem Env:") {
			t.Fatalf("shellPathCommandOutput args = %#v, want PowerShell env probe", args)
		}
		output := strings.Join([]string{
			"profile noise",
			shellEnvStart,
			"OPENAI_API_KEY=sk-from-pwsh",
			"CODEX_HOME=C:\\Users\\demo\\.codex-work",
			shellEnvEnd,
		}, "\n")
		return []byte(output), nil
	}

	env, ok := resolveShellEnvFromCommands(shellPathTimeout, windowsShellEnvCommands()[:1])
	if !ok {
		t.Fatal("resolveShellEnvFromCommands() ok = false, want true")
	}
	if got := env["OPENAI_API_KEY"]; got != "sk-from-pwsh" {
		t.Fatalf("OPENAI_API_KEY = %q, want sk-from-pwsh", got)
	}
	if got := env["CODEX_HOME"]; got != `C:\Users\demo\.codex-work` {
		t.Fatalf("CODEX_HOME = %q, want custom Codex home", got)
	}
}

func TestShellPathCommandsCoverPlatforms(t *testing.T) {
	unixCommands := shellPathCommandsForGOOS("darwin", "/custom/fish")
	if len(unixCommands) == 0 {
		t.Fatal("expected unix shell commands")
	}
	if unixCommands[0].command != "/custom/fish" {
		t.Fatalf("first unix command = %q, want user's shell", unixCommands[0].command)
	}
	if !strings.Contains(unixCommands[0].args[len(unixCommands[0].args)-1], "string join : $PATH") {
		t.Fatalf("fish command did not use fish PATH syntax: %#v", unixCommands[0].args)
	}

	envCommands := shellEnvCommandsForGOOS("darwin", "/custom/fish")
	if len(envCommands) == 0 {
		t.Fatal("expected shell env commands")
	}
	if envCommands[0].command != "/custom/fish" {
		t.Fatalf("first shell env command = %q, want user's shell", envCommands[0].command)
	}
	if !strings.Contains(envCommands[0].args[len(envCommands[0].args)-1], shellEnvStart) {
		t.Fatalf("shell env command did not include env sentinel: %#v", envCommands[0].args)
	}

	windowsCommands := shellPathCommandsForGOOS("windows", "")
	names := make([]string, 0, len(windowsCommands))
	for _, cmd := range windowsCommands {
		names = append(names, cmd.command)
	}
	joined := strings.Join(names, ",")
	for _, name := range []string{"pwsh.exe", "powershell.exe", "cmd.exe"} {
		if !strings.Contains(joined, name) {
			t.Fatalf("windows shell commands = %v, missing %s", names, name)
		}
	}

	windowsEnvCommands := shellEnvCommandsForGOOS("windows", "")
	names = names[:0]
	for _, cmd := range windowsEnvCommands {
		names = append(names, cmd.command)
		if len(cmd.args) == 0 || !strings.Contains(cmd.args[len(cmd.args)-1], shellEnvStart) {
			t.Fatalf("windows env command %q did not include env sentinel: %#v", cmd.command, cmd.args)
		}
	}
	joined = strings.Join(names, ",")
	for _, name := range []string{"pwsh.exe", "powershell.exe"} {
		if !strings.Contains(joined, name) {
			t.Fatalf("windows shell env commands = %v, missing %s", names, name)
		}
	}
}

func TestLoadEnvFromShellProfilesDoesNotOverwriteShellSnapshot(t *testing.T) {
	dir := t.TempDir()
	t.Setenv("HOME", dir)
	runtimeconfig.SetShellEnvForRuntimeResolution(map[string]string{
		"LLMPROXY_OAI_KEY": "from-shell",
	})
	t.Cleanup(func() {
		runtimeconfig.SetShellEnvForRuntimeResolution(nil)
	})

	profilePath := filepath.Join(dir, ".zshenv")
	content := "export LLMPROXY_OAI_KEY=from-profile\n"
	if err := os.WriteFile(profilePath, []byte(content), 0644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}

	LoadEnvFromShellProfiles()

	env := runtimeconfig.ShellEnvForRuntimeResolution()
	if got := env["LLMPROXY_OAI_KEY"]; got != "from-shell" {
		t.Errorf("runtime env = %q, want from-shell", got)
	}
}

func TestParseProfileEnvsNonexistent(t *testing.T) {
	env, err := parseProfileEnvs("/nonexistent/.bashrc")
	if err == nil {
		t.Error("expected error for nonexistent file")
	}
	if len(env) != 0 {
		t.Errorf("env = %#v, want empty", env)
	}
}

func TestShellProfiles(t *testing.T) {
	profiles := shellProfiles("/home/user")
	if runtime.GOOS == "windows" {
		if len(profiles) != 0 {
			t.Errorf("expected no profiles on windows, got %d", len(profiles))
		}
		return
	}
	if len(profiles) == 0 {
		t.Error("expected at least one profile")
	}

	// All paths should be under the given homeDir.
	for _, p := range profiles {
		if !strings.HasPrefix(p, "/home/user/") {
			t.Errorf("profile path %q should be under /home/user/", p)
		}
	}
}

func TestShellProfilesIncludesMacOSZProfile(t *testing.T) {
	profiles := shellProfilesForGOOS("darwin", "/Users/alice")
	want := filepath.Join("/Users/alice", ".zprofile")
	for _, profile := range profiles {
		if profile == want {
			return
		}
	}
	t.Fatalf("shellProfilesForGOOS(darwin) = %v, want %q included", profiles, want)
}