package agent

import (
	"sort"
	"strconv"
	"strings"

	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// extractBearerToken returns the bearer token from an MCP headers map, or ""
// if no Authorization header is present.
func extractBearerToken(headers map[string]string) string {
	for k, v := range headers {
		value := strings.TrimSpace(v)
		if strings.EqualFold(strings.TrimSpace(k), "Authorization") && strings.HasPrefix(strings.ToLower(value), "bearer ") {
			return strings.TrimSpace(value[len("bearer "):])
		}
	}
	return ""
}

// mcpAuthTokenEnvVarForServer returns the environment variable name used to
// pass a bearer token for the given MCP server name.
func mcpAuthTokenEnvVarForServer(name string) string {
	return "AGENTTEAM_MCP_AUTH_TOKEN_" + normalizeMCPEnvNamePart(name)
}

func mcpEnvHTTPHeadersForServer(name string, headers map[string]string) map[string]string {
	if len(headers) == 0 {
		return nil
	}
	result := make(map[string]string, len(headers))
	for rawHeader, rawValue := range headers {
		header := strings.TrimSpace(rawHeader)
		if header == "" || strings.EqualFold(header, "Authorization") || strings.TrimSpace(rawValue) == "" {
			continue
		}
		result[header] = mcpHTTPHeaderEnvVarForServer(name, header)
	}
	if len(result) == 0 {
		return nil
	}
	return result
}

func applyMCPHeaderEnv(env map[string]string, mcps []types.NamedMCPConfig) map[string]string {
	for _, m := range mcps {
		if token := extractBearerToken(m.Headers); token != "" {
			if env == nil {
				env = make(map[string]string)
			}
			env[mcpAuthTokenEnvVarForServer(m.Name)] = token
		}
		for rawHeader, rawValue := range m.Headers {
			header := strings.TrimSpace(rawHeader)
			value := strings.TrimSpace(rawValue)
			if header == "" || strings.EqualFold(header, "Authorization") || value == "" {
				continue
			}
			if env == nil {
				env = make(map[string]string)
			}
			env[mcpHTTPHeaderEnvVarForServer(m.Name, header)] = value
		}
	}
	return env
}

func mcpHTTPHeaderEnvVarForServer(name string, header string) string {
	return "AGENTTEAM_MCP_HTTP_HEADER_" + normalizeMCPEnvNamePart(name) + "_" + normalizeMCPEnvNamePart(header)
}

func mcpStdioEnvVarForServer(name string, envName string) string {
	return "AGENTTEAM_MCP_STDIO_ENV_" + normalizeMCPEnvNamePart(name) + "_" + normalizeMCPEnvNamePart(envName)
}

func tomlInlineStringMap(values map[string]string) string {
	if len(values) == 0 {
		return "{}"
	}
	keys := make([]string, 0, len(values))
	for key := range values {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	parts := make([]string, 0, len(keys))
	for _, key := range keys {
		parts = append(parts, strconv.Quote(key)+"="+strconv.Quote(values[key]))
	}
	return "{" + strings.Join(parts, ", ") + "}"
}

func normalizeMCPEnvNamePart(raw string) string {
	var b strings.Builder
	lastUnderscore := false
	for _, r := range strings.TrimSpace(raw) {
		switch {
		case r >= 'a' && r <= 'z':
			b.WriteRune(r - 'a' + 'A')
			lastUnderscore = false
		case (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9'):
			b.WriteRune(r)
			lastUnderscore = false
		default:
			if b.Len() == 0 || lastUnderscore {
				continue
			}
			b.WriteByte('_')
			lastUnderscore = true
		}
	}
	result := strings.Trim(b.String(), "_")
	if result == "" {
		return "VALUE"
	}
	return result
}