package agent

import (
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

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

func TestApplyCodexProviderConfigAddsOverridesAndEnv(t *testing.T) {
	t.Parallel()

	args, env := applyCodexProviderConfig([]string{"exec"}, nil, &drivercontract.ProviderConfig{
		ID:       "direct-upstream",
		Name:     "Direct Upstream",
		BaseURL:  "https://llm.example.com/v1",
		APIKey:   "agent-token",
		WireAPIs: []string{provider.WireAPIResponses},
	})

	require.Contains(t, args, "-c")
	require.Contains(t, args, `model_provider="direct-upstream"`)
	require.Contains(t, args, `model_providers.direct-upstream.name="Direct Upstream"`)
	require.Contains(t, args, `model_providers.direct-upstream.base_url="https://llm.example.com/v1"`)
	require.Contains(t, args, `model_providers.direct-upstream.wire_api="responses"`)
	require.Contains(t, args, `model_providers.direct-upstream.env_key="AGENTTEAM_CODEX_PROVIDER_API_KEY"`)
	require.Equal(t, "agent-token", env[codexProviderAPIKeyEnvVar])
}

func TestApplyCodexProviderConfigPreservesEnvKeyWhenAPIKeyIsShellBacked(t *testing.T) {
	t.Parallel()

	args, env := applyCodexProviderConfig([]string{"exec"}, nil, &drivercontract.ProviderConfig{
		ID:       "runtime-default",
		Name:     "Runtime Default",
		BaseURL:  "https://llm.example.com/v1",
		EnvKey:   "OPENAI_API_KEY",
		WireAPIs: []string{provider.WireAPIResponses},
	})

	require.Contains(t, args, `model_provider="runtime-default"`)
	require.Contains(t, args, `model_providers.runtime-default.env_key="OPENAI_API_KEY"`)
	require.NotContains(t, args, `model_providers.runtime-default.env_key="AGENTTEAM_CODEX_PROVIDER_API_KEY"`)
	require.Nil(t, env)
}

func TestApplyCodexProviderConfigPrefersResponsesWireAPI(t *testing.T) {
	t.Parallel()

	args, _ := applyCodexProviderConfig([]string{"exec"}, nil, &drivercontract.ProviderConfig{
		ID:       "direct-upstream",
		Name:     "Direct Upstream",
		BaseURL:  "https://llm.example.com/v1",
		WireAPIs: []string{provider.WireAIOpenAI, provider.WireAPIResponses},
	})

	require.Contains(t, args, `model_providers.direct-upstream.wire_api="responses"`)
	require.NotContains(t, args, `model_providers.direct-upstream.wire_api="openai"`)
}

func TestApplyCodexProviderConfigDefaultsToResponsesWireAPI(t *testing.T) {
	t.Parallel()

	args, _ := applyCodexProviderConfig([]string{"exec"}, nil, &drivercontract.ProviderConfig{
		ID:      "direct-upstream",
		Name:    "Direct Upstream",
		BaseURL: "https://llm.example.com/v1",
	})

	require.Contains(t, args, `model_providers.direct-upstream.wire_api="responses"`)
	require.NotContains(t, args, `model_providers.direct-upstream.wire_api="openai"`)
}

func TestApplyCodexProviderConfigWithoutProviderIsNoop(t *testing.T) {
	t.Parallel()

	baseArgs := []string{"exec"}
	baseEnv := map[string]string{"EXISTING": "1"}

	args, env := applyCodexProviderConfig(baseArgs, baseEnv, nil)

	require.Equal(t, baseArgs, args)
	require.Equal(t, baseEnv, env)
}

func TestApplyCodexProviderConfigWithEmptyProviderFieldsIsNoop(t *testing.T) {
	t.Parallel()

	baseArgs := []string{"exec"}
	baseEnv := map[string]string{"EXISTING": "1"}

	args, env := applyCodexProviderConfig(baseArgs, baseEnv, &drivercontract.ProviderConfig{})

	require.Equal(t, baseArgs, args)
	require.Equal(t, baseEnv, env)
}

func TestBuildRuntimeEnvWithProviderStripsAmbientVendorEnv(t *testing.T) {
	runtimeconfig.SetShellEnvForRuntimeResolution(nil)
	t.Cleanup(func() {
		runtimeconfig.SetShellEnvForRuntimeResolution(nil)
	})
	t.Setenv("OPENAI_API_KEY", "sk-upstream")
	t.Setenv("GLM_API_KEY", "sk-glm")

	env := buildRuntimeEnvWithProvider(map[string]string{
		codexProviderAPIKeyEnvVar: "sk-direct",
		"OTHER_ENV":               "1",
	}, &drivercontract.ProviderConfig{
		BaseURL: "https://llm.example.com/v1",
	})

	joined := strings.Join(env, "\n")
	require.NotContains(t, joined, "OPENAI_API_KEY=sk-upstream")
	require.NotContains(t, joined, "GLM_API_KEY=sk-glm")
	require.Contains(t, joined, codexProviderAPIKeyEnvVar+"=sk-direct")
	require.Contains(t, joined, "OTHER_ENV=1")
}

func TestBuildRuntimeEnvWithProviderKeepsConfiguredEnvKeyWhenShellBacked(t *testing.T) {
	runtimeconfig.SetShellEnvForRuntimeResolution(map[string]string{
		"OPENAI_API_KEY":    "sk-shell",
		"GLM_API_KEY":       "sk-glm",
		"CUSTOM_LLM_REGION": "us-east-1",
	})
	t.Cleanup(func() {
		runtimeconfig.SetShellEnvForRuntimeResolution(nil)
	})

	env := buildRuntimeEnvWithProvider(nil, &drivercontract.ProviderConfig{
		BaseURL: "https://llm.example.com/v1",
		EnvKey:  "OPENAI_API_KEY",
	})

	joined := strings.Join(env, "\n")
	require.Contains(t, joined, "OPENAI_API_KEY=sk-shell")
	require.NotContains(t, joined, "GLM_API_KEY=sk-glm")
	require.Contains(t, joined, "CUSTOM_LLM_REGION=us-east-1")
}

func TestBuildRuntimeEnvWithProviderReplacesAmbientExecCredential(t *testing.T) {
	runtimeconfig.SetShellEnvForRuntimeResolution(map[string]string{
		"AGENTTEAM_EXEC_BASE_URL":   "http://127.0.0.1:1",
		"AGENTTEAM_EXEC_TOKEN":      "shell-bearer",
		"AGENTTEAM_EXEC_TOKEN_FILE": "/private/shell-token",
	})
	t.Cleanup(func() {
		runtimeconfig.SetShellEnvForRuntimeResolution(nil)
	})
	t.Setenv("AGENTTEAM_EXEC_BASE_URL", "http://127.0.0.1:2")
	t.Setenv("AGENTTEAM_EXEC_TOKEN", "ambient-bearer")
	t.Setenv("AGENTTEAM_EXEC_TOKEN_FILE", "/private/ambient-token")

	env := buildRuntimeEnvWithProvider(map[string]string{
		"AGENTTEAM_EXEC_BASE_URL":   "http://127.0.0.1:43123",
		"AGENTTEAM_EXEC_TOKEN":      "stale-runtime-bearer",
		"AGENTTEAM_EXEC_TOKEN_FILE": "/runtime/current-token",
		"agentteam_exec_base_url":   "http://127.0.0.1:3",
		"agentteam_exec_token":      "mixed-case-bearer",
		"agentteam_exec_token_file": "/private/mixed-case-token",
	}, nil)

	require.Equal(t, []string{"http://127.0.0.1:43123"}, runtimeEnvValues(env, "AGENTTEAM_EXEC_BASE_URL"))
	require.Empty(t, runtimeEnvValues(env, "AGENTTEAM_EXEC_TOKEN"))
	require.Equal(t, []string{"/runtime/current-token"}, runtimeEnvValues(env, "AGENTTEAM_EXEC_TOKEN_FILE"))
}

func runtimeEnvValues(env []string, wantKey string) []string {
	values := make([]string, 0, 1)
	for _, entry := range env {
		key, value, ok := strings.Cut(entry, "=")
		if ok && strings.EqualFold(key, wantKey) {
			values = append(values, value)
		}
	}
	return values
}