package agent

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"

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

func buildRuntimeEnv(extra map[string]string) []string {
	return buildRuntimeEnvWithProvider(extra, nil)
}

func buildRuntimeEnvWithProvider(extra map[string]string, prov *drivercontract.ProviderConfig) []string {
	env := append([]string(nil), os.Environ()...)
	// Filter out CLAUDECODE env var — when Claude is running as a daemon-managed
	// agent, it should not know it's being driven by another Claude instance.
	filtered := env[:0]
	for _, e := range env {
		if strings.HasPrefix(e, "CLAUDECODE=") {
			continue
		}
		if providerUpstreamConfigured(prov) && shouldStripAmbientProviderEnvEntry(e, prov) {
			continue
		}
		filtered = append(filtered, e)
	}
	env = filtered
	env = appendRuntimeShellEnvWithProviderFilter(env, prov)
	env = append(env, "FORCE_COLOR=0", "NO_COLOR=1")
	return appendRuntimeExtraEnv(env, extra)
}

// buildClaudeEnv constructs the environment for the Claude CLI.
// When running as root (euid == 0), it forces IS_SANDBOX=1 to prevent
// Claude Code from refusing to run in dangerous-permissions mode.
// If a provider config is supplied, it injects ANTHROPIC_BASE_URL and the
// Anthropic credential variable expected by the configured upstream.
func buildClaudeEnv(extra map[string]string, prov *drivercontract.ProviderConfig, smallModel string) []string {
	isRoot := os.Geteuid() == 0
	env := append([]string(nil), os.Environ()...)
	filtered := env[:0]
	for _, e := range env {
		if strings.HasPrefix(e, "CLAUDECODE=") {
			continue
		}
		if isRoot && strings.HasPrefix(e, "IS_SANDBOX=") {
			continue
		}
		if providerUpstreamConfigured(prov) && shouldStripAmbientProviderEnvEntry(e, prov) {
			continue
		}
		filtered = append(filtered, e)
	}
	env = filtered
	env = appendRuntimeShellEnvWithProviderFilter(env, prov)
	env = append(env, "FORCE_COLOR=0", "NO_COLOR=1")
	if isRoot {
		env = append(env, "IS_SANDBOX=1")
	}
	// Inject provider credentials so Claude uses the configured upstream endpoint.
	if baseURL := claudeAnthropicBaseURL(prov); baseURL != "" {
		env = append(env, "ANTHROPIC_BASE_URL="+baseURL)
		if strings.TrimSpace(prov.APIKey) != "" {
			envName := claudeProviderCredentialEnvName(prov, baseURL)
			env = append(env, envName+"="+strings.TrimSpace(prov.APIKey))
		}
	}
	if strings.TrimSpace(smallModel) != "" {
		env = append(env, "CLAUDE_SMALL_MODEL="+strings.TrimSpace(smallModel))
	}
	return appendRuntimeExtraEnv(env, extra)
}

func appendRuntimeShellEnvWithProviderFilter(env []string, prov *drivercontract.ProviderConfig) []string {
	return runtimeconfig.AppendShellEnvForRuntimeIf(env, func(key, value string) bool {
		if execruntimeenv.IsDaemonOwnedKey(key) {
			return false
		}
		return !providerUpstreamConfigured(prov) || !shouldStripAmbientProviderEnvKey(key, prov)
	})
}

func appendRuntimeExtraEnv(env []string, extra map[string]string) []string {
	baseURL, hasBaseURL := extra[execruntimeenv.BaseURL]
	tokenFile, hasTokenFile := extra[execruntimeenv.TokenFile]
	for key, value := range extra {
		// Ignore every spelling of the daemon-owned names so a mixed-case
		// per-runtime override cannot shadow them on Windows. Re-append only the
		// resolved base/file values authored by the dispatcher; the legacy raw
		// bearer is never inherited by a newly launched runtime.
		if execruntimeenv.IsDaemonOwnedKey(key) {
			continue
		}
		env = append(env, fmt.Sprintf("%s=%s", key, value))
	}
	if hasBaseURL {
		env = append(env, execruntimeenv.BaseURL+"="+baseURL)
	}
	if hasTokenFile {
		env = append(env, execruntimeenv.TokenFile+"="+tokenFile)
	}
	return env
}

func buildClaudeSettingsOverrides(cfg drivercontract.LaunchConfig) string {
	baseURL := claudeAnthropicBaseURL(cfg.Provider)
	if baseURL == "" {
		return ""
	}

	env := map[string]string{
		"ANTHROPIC_BASE_URL": baseURL,
	}
	if apiKey := strings.TrimSpace(cfg.Provider.APIKey); apiKey != "" {
		switch claudeProviderCredentialEnvName(cfg.Provider, baseURL) {
		case "ANTHROPIC_AUTH_TOKEN":
			env["ANTHROPIC_AUTH_TOKEN"] = apiKey
		default:
			env["ANTHROPIC_API_KEY"] = apiKey
		}
	}
	if model := strings.TrimSpace(cfg.Model); model != "" {
		env["ANTHROPIC_MODEL"] = model
	}
	if smallModel := strings.TrimSpace(cfg.SmallModel); smallModel != "" {
		env["ANTHROPIC_SMALL_FAST_MODEL"] = smallModel
	}

	payload := map[string]any{"env": env}
	data, err := json.Marshal(payload)
	if err != nil {
		return ""
	}
	return string(data)
}

func claudeProviderCredentialEnvName(prov *drivercontract.ProviderConfig, baseURL string) string {
	if prov != nil {
		switch strings.TrimSpace(prov.EnvKey) {
		case "ANTHROPIC_AUTH_TOKEN":
			return "ANTHROPIC_AUTH_TOKEN"
		case "ANTHROPIC_API_KEY":
			return "ANTHROPIC_API_KEY"
		}
	}
	return claudeAnthropicAPIKeyEnvName(baseURL)
}

func claudeAnthropicAPIKeyEnvName(baseURL string) string {
	normalized := strings.TrimRight(strings.TrimSpace(baseURL), "/")
	if _, ok := claudeAuthTokenBaseURLs[normalized]; ok {
		return "ANTHROPIC_AUTH_TOKEN"
	}
	return "ANTHROPIC_API_KEY"
}

func providerUpstreamConfigured(prov *drivercontract.ProviderConfig) bool {
	return prov != nil && strings.TrimSpace(prov.BaseURL) != ""
}

func claudeAnthropicBaseURL(prov *drivercontract.ProviderConfig) string {
	if prov == nil {
		return ""
	}
	baseURL := strings.TrimSpace(prov.BaseURL)
	if baseURL == "" {
		return ""
	}
	baseURL = strings.TrimRight(baseURL, "/")
	switch {
	case strings.HasSuffix(baseURL, "/v1/messages"):
		return strings.TrimSuffix(baseURL, "/v1/messages")
	case strings.HasSuffix(baseURL, "/messages"):
		return strings.TrimSuffix(baseURL, "/messages")
	case strings.HasSuffix(baseURL, "/v1"):
		return strings.TrimSuffix(baseURL, "/v1")
	default:
		return baseURL
	}
}

var ambientProviderEnvKeys = map[string]struct{}{
	"ANTHROPIC_API_KEY":    {},
	"ANTHROPIC_AUTH_TOKEN": {},
	"ANTHROPIC_BASE_URL":   {},
	"ARK_API_KEY":          {},
	"ARK_BASE_URL":         {},
	"BAILIAN_API_KEY":      {},
	"BAILIAN_BASE_URL":     {},
	"DASHSCOPE_API_KEY":    {},
	"DASHSCOPE_BASE_URL":   {},
	"DEEPSEEK_API_KEY":     {},
	"DEEPSEEK_BASE_URL":    {},
	"GEMINI_API_KEY":       {},
	"GEMINI_BASE_URL":      {},
	"GLM_API_KEY":          {},
	"GLM_BASE_URL":         {},
	"GOOGLE_API_KEY":       {},
	"GOOGLE_BASE_URL":      {},
	"MOONSHOT_API_KEY":     {},
	"MOONSHOT_BASE_URL":    {},
	"OPENAI_API_KEY":       {},
	"OPENAI_BASE_URL":      {},
	"OPENROUTER_API_KEY":   {},
	"OPENROUTER_BASE_URL":  {},
	"XAI_API_KEY":          {},
	"XAI_BASE_URL":         {},
}

func shouldStripAmbientProviderEnvEntry(entry string, prov *drivercontract.ProviderConfig) bool {
	key := entry
	if idx := strings.IndexByte(entry, '='); idx >= 0 {
		key = entry[:idx]
	}
	return shouldStripAmbientProviderEnvKey(key, prov)
}

func shouldStripAmbientProviderEnvKey(key string, prov *drivercontract.ProviderConfig) bool {
	if prov != nil && strings.TrimSpace(prov.APIKey) == "" {
		if envKey := strings.TrimSpace(prov.EnvKey); envKey != "" && strings.EqualFold(strings.TrimSpace(key), envKey) {
			return false
		}
	}
	_, ok := ambientProviderEnvKeys[strings.TrimSpace(key)]
	return ok
}