package agent

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"github.com/pelletier/go-toml/v2"

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

var writeKimiACPProviderConfigFile = os.WriteFile
var kimiACPConfigMu sync.Mutex

type kimiACPPreparedLaunch struct {
	Args       []string
	Env        []string
	WorkDir    string
	ConfigPath string
	Cleanup    func()
}

func prepareKimiACPLaunch(cfg drivercontract.LaunchConfig) (kimiACPPreparedLaunch, error) {
	workDir, err := filepath.Abs(strings.TrimSpace(cfg.WorkDir))
	if err != nil || strings.TrimSpace(cfg.WorkDir) == "" {
		return kimiACPPreparedLaunch{}, fmt.Errorf("kimi: valid workdir is required")
	}
	workDir = filepath.Clean(workDir)
	if err := os.MkdirAll(workDir, 0o755); err != nil {
		return kimiACPPreparedLaunch{}, fmt.Errorf("kimi: create workdir: %w", err)
	}
	if _, err := cfg.EffectivePermissionMode(); err != nil {
		return kimiACPPreparedLaunch{}, err
	}
	if cfg.DisableNativeTools {
		return kimiACPPreparedLaunch{}, fmt.Errorf("kimi: ACP cannot enforce DisableNativeTools; refusing to expose Kimi native tools")
	}

	env := runtimeconfig.RuntimeEnvOverrides(cfg.Env, cfg.MCPs)
	if env == nil {
		env = make(map[string]string)
	}
	prepared := kimiACPPreparedLaunch{
		Args:    []string{"acp"},
		WorkDir: workDir,
		Cleanup: func() {},
	}
	if cfg.Provider == nil {
		prepared.Env = buildRuntimeEnv(env)
		return prepared, nil
	}

	removeKimiACPProviderEnv(env)
	kimiHome, err := kimiACPHome(cfg)
	if err != nil {
		return kimiACPPreparedLaunch{}, err
	}
	env["KIMI_CODE_HOME"] = kimiHome
	configPath := filepath.Join(kimiHome, "config.toml")
	kimiACPConfigMu.Lock()
	defer kimiACPConfigMu.Unlock()
	// A prior crash may leave a launch-owned provider config behind. It is
	// safe to replace only inside AgentTeam's per-agent runtime directory.
	if err := os.Remove(configPath); err != nil && !os.IsNotExist(err) {
		return kimiACPPreparedLaunch{}, fmt.Errorf("kimi: remove stale provider config: %w", err)
	}
	useEnv, err := prepareKimiACPProviderEnv(env, cfg)
	if err != nil {
		return kimiACPPreparedLaunch{}, err
	}
	if !useEnv {
		ownerToken, err := newKimiACPConfigOwnerToken()
		if err != nil {
			return kimiACPPreparedLaunch{}, err
		}
		if err := writeKimiACPProviderConfig(configPath, cfg, ownerToken); err != nil {
			return kimiACPPreparedLaunch{}, err
		}
		prepared.ConfigPath = configPath
		prepared.Cleanup = func() { removeKimiACPConfigIfOwned(configPath, ownerToken) }
	}
	prepared.Env = sanitizeKimiACPProviderProcessEnv(buildRuntimeEnvWithProvider(env, cfg.Provider), useEnv)
	return prepared, nil
}

func removeKimiACPProviderEnv(env map[string]string) {
	for key := range env {
		normalizedKey := strings.ToUpper(strings.TrimSpace(key))
		if normalizedKey == "KIMI_CODE_HOME" || strings.HasPrefix(normalizedKey, "KIMI_MODEL_") {
			delete(env, key)
		}
	}
}

func sanitizeKimiACPProviderProcessEnv(env []string, useModelEnv bool) []string {
	seen := make(map[string]struct{})
	filtered := make([]string, 0, len(env))
	for index := len(env) - 1; index >= 0; index-- {
		entry := env[index]
		key, _, _ := strings.Cut(entry, "=")
		normalizedKey := strings.ToUpper(strings.TrimSpace(key))
		if strings.HasPrefix(normalizedKey, "KIMI_MODEL_") && !useModelEnv {
			continue
		}
		if normalizedKey == "KIMI_CODE_HOME" || strings.HasPrefix(normalizedKey, "KIMI_MODEL_") {
			if _, exists := seen[normalizedKey]; exists {
				continue
			}
			seen[normalizedKey] = struct{}{}
		}
		filtered = append(filtered, entry)
	}
	for left, right := 0, len(filtered)-1; left < right; left, right = left+1, right-1 {
		filtered[left], filtered[right] = filtered[right], filtered[left]
	}
	return filtered
}

func newKimiACPConfigOwnerToken() (string, error) {
	var token [16]byte
	if _, err := rand.Read(token[:]); err != nil {
		return "", fmt.Errorf("kimi: generate provider config owner: %w", err)
	}
	return hex.EncodeToString(token[:]), nil
}

func removeKimiACPConfigIfOwned(configPath, ownerToken string) {
	kimiACPConfigMu.Lock()
	defer kimiACPConfigMu.Unlock()
	payload, err := os.ReadFile(configPath)
	if err != nil || !strings.HasPrefix(string(payload), "# agentteam-owner: "+ownerToken+"\n") {
		return
	}
	_ = os.Remove(configPath)
}

func kimiACPHome(cfg drivercontract.LaunchConfig) (string, error) {
	root := strings.TrimSpace(cfg.RuntimeStateDir)
	if root == "" {
		return "", fmt.Errorf("kimi: RuntimeStateDir is required for an injected provider")
	}
	abs, err := filepath.Abs(root)
	if err != nil {
		return "", fmt.Errorf("kimi: resolve runtime state directory: %w", err)
	}
	home := filepath.Join(filepath.Clean(abs), "kimi-code")
	if err := ensureKimiPrivateDirectory(home); err != nil {
		return "", fmt.Errorf("kimi: create private ACP home: %w", err)
	}
	return home, nil
}

func prepareKimiACPProviderEnv(env map[string]string, cfg drivercontract.LaunchConfig) (bool, error) {
	prov := cfg.Provider
	apiKey, model, err := kimiACPProviderFields(cfg)
	if err != nil {
		return false, err
	}
	providerType, err := kimiACPProviderType(prov)
	if err != nil {
		return false, err
	}
	if providerType == "openai_responses" || len(prov.Headers) > 0 {
		return false, nil
	}
	env["KIMI_MODEL_NAME"] = model
	env["KIMI_MODEL_DISPLAY_NAME"] = firstNonEmptyString(strings.TrimSpace(prov.Name), model)
	env["KIMI_MODEL_API_KEY"] = apiKey
	env["KIMI_MODEL_BASE_URL"] = strings.TrimSpace(prov.BaseURL)
	env["KIMI_MODEL_PROVIDER_TYPE"] = providerType
	return true, nil
}

func writeKimiACPProviderConfig(path string, cfg drivercontract.LaunchConfig, ownerToken string) error {
	prov := cfg.Provider
	apiKey, model, err := kimiACPProviderFields(cfg)
	if err != nil {
		return err
	}
	providerType, err := kimiACPProviderType(prov)
	if err != nil {
		return err
	}
	alias := "agentteam-runtime"
	config := map[string]any{
		"default_model": alias,
		"providers": map[string]any{
			alias: map[string]any{
				"type":           providerType,
				"base_url":       strings.TrimSpace(prov.BaseURL),
				"api_key":        apiKey,
				"custom_headers": runtimeconfig.CloneStringMap(prov.Headers),
			},
		},
		"models": map[string]any{
			alias: map[string]any{
				"provider":         alias,
				"model":            model,
				"max_context_size": 262144,
				"display_name":     firstNonEmptyString(strings.TrimSpace(prov.Name), model),
			},
		},
	}
	payload, err := toml.Marshal(config)
	if err != nil {
		return fmt.Errorf("kimi: encode provider config: %w", err)
	}
	payload = append([]byte("# agentteam-owner: "+ownerToken+"\n"), payload...)
	if err := writeKimiACPProviderConfigFile(path, payload, 0o600); err != nil {
		return fmt.Errorf("kimi: write provider config: %w", err)
	}
	return nil
}

func kimiACPProviderFields(cfg drivercontract.LaunchConfig) (apiKey string, model string, err error) {
	prov := cfg.Provider
	if prov == nil {
		return "", "", fmt.Errorf("kimi: provider config is required")
	}
	apiKey = strings.TrimSpace(prov.APIKey)
	if apiKey == "" && strings.TrimSpace(prov.EnvKey) != "" {
		if value, ok := runtimeconfig.RuntimeEnvValue(strings.TrimSpace(prov.EnvKey)); ok {
			apiKey = strings.TrimSpace(value)
		}
	}
	model = normalizeKimiModelName(cfg.Model, prov)
	if apiKey == "" || model == "" || strings.TrimSpace(prov.BaseURL) == "" {
		return "", "", fmt.Errorf("kimi: provider, model, base URL, and API key are required for an explicit provider override")
	}
	return apiKey, model, nil
}

func kimiACPProviderType(prov *drivercontract.ProviderConfig) (string, error) {
	if prov == nil {
		return "", fmt.Errorf("kimi: provider is required")
	}
	for _, wireAPI := range prov.WireAPIs {
		switch strings.TrimSpace(wireAPI) {
		case provider.WireAPIAnthropic:
			return "anthropic", nil
		case provider.WireAIOpenAI:
			return "openai", nil
		case provider.WireAPIResponses:
			return "openai_responses", nil
		}
	}
	return "openai", nil
}

func kimiACPPromptBlocks(input drivercontract.RuntimeInput) ([]map[string]any, error) {
	blocks := make([]map[string]any, 0, len(input.Attachments)+1)
	if strings.TrimSpace(input.PromptText()) != "" {
		blocks = append(blocks, map[string]any{"type": "text", "text": input.PromptText()})
	}
	for _, attachment := range input.Attachments {
		mime := strings.ToLower(strings.TrimSpace(attachment.MIME))
		if strings.HasPrefix(mime, "image/") {
			path := runtimeattachments.LocalPath(attachment)
			if strings.TrimSpace(path) == "" {
				return nil, fmt.Errorf("kimi: image attachment %q has no local path", attachment.Filename)
			}
			data, err := os.ReadFile(path)
			if err != nil {
				return nil, fmt.Errorf("kimi: read image attachment %q: %w", attachment.Filename, err)
			}
			blocks = append(blocks, map[string]any{
				"type":     "image",
				"mimeType": attachment.MIME,
				"data":     base64.StdEncoding.EncodeToString(data),
			})
			continue
		}
		if uri := strings.TrimSpace(attachment.URI); uri != "" {
			blocks = append(blocks, map[string]any{
				"type": "resource_link",
				"uri":  uri,
				"name": firstNonEmptyString(strings.TrimSpace(attachment.Filename), uri),
			})
		}
	}
	if len(blocks) == 0 {
		blocks = append(blocks, map[string]any{"type": "text", "text": ""})
	}
	return blocks, nil
}

func kimiACPMCPServers(mcps []types.NamedMCPConfig) ([]map[string]any, error) {
	servers := make([]map[string]any, 0, len(mcps))
	for _, mcp := range mcps {
		name := strings.TrimSpace(mcp.Name)
		if name == "" {
			continue
		}
		switch mcp.EffectiveTransport() {
		case types.MCPTransportStdio:
			if strings.TrimSpace(mcp.CWD) != "" {
				return nil, fmt.Errorf("kimi: ACP stdio MCP %q cannot preserve an explicit cwd", name)
			}
			command := strings.TrimSpace(mcp.Command)
			if command == "" {
				return nil, fmt.Errorf("kimi: ACP stdio MCP %q requires a command", name)
			}
			env := make([]map[string]string, 0, len(mcp.Env))
			keys := make([]string, 0, len(mcp.Env))
			for key := range mcp.Env {
				keys = append(keys, key)
			}
			sort.Strings(keys)
			for _, key := range keys {
				env = append(env, map[string]string{"name": key, "value": mcp.Env[key]})
			}
			servers = append(servers, map[string]any{
				"name": name, "command": command,
				"args": append([]string(nil), mcp.Args...), "env": env,
			})
		default:
			url := strings.TrimSpace(mcp.URL)
			if url == "" {
				return nil, fmt.Errorf("kimi: ACP HTTP MCP %q requires a URL", name)
			}
			headers := make([]map[string]string, 0, len(mcp.Headers))
			keys := make([]string, 0, len(mcp.Headers))
			for key := range mcp.Headers {
				keys = append(keys, key)
			}
			sort.Strings(keys)
			for _, key := range keys {
				headers = append(headers, map[string]string{"name": key, "value": mcp.Headers[key]})
			}
			servers = append(servers, map[string]any{
				"type": "http", "name": name, "url": url, "headers": headers,
			})
		}
	}
	return servers, nil
}