package agent

import (
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"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/promptcompiler"
	"github.com/tidbcloud/AgentTeam/desktop/internal/provider"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

type kimiPreparedLaunch struct {
	Args             []string
	Env              []string
	WorkDir          string
	ShareDir         string
	AgentFile        string
	SystemPromptFile string
	MCPConfigFile    string
	Cleanup          func()
}

var writeKimiProviderConfigFile = os.WriteFile

func (prepared *kimiPreparedLaunch) addCleanup(cleanup func()) {
	if prepared == nil || cleanup == nil {
		return
	}
	previous := prepared.Cleanup
	prepared.Cleanup = func() {
		cleanup()
		if previous != nil {
			previous()
		}
	}
}

func prepareKimiLaunch(cfg drivercontract.LaunchConfig, sessionID string) (kimiPreparedLaunch, error) {
	workDir, err := filepath.Abs(strings.TrimSpace(cfg.WorkDir))
	if err != nil || strings.TrimSpace(cfg.WorkDir) == "" {
		return kimiPreparedLaunch{}, fmt.Errorf("kimi: valid workdir is required")
	}
	workDir = filepath.Clean(workDir)
	if err := os.MkdirAll(workDir, 0o755); err != nil {
		return kimiPreparedLaunch{}, fmt.Errorf("kimi: create workdir: %w", err)
	}

	env := runtimeconfig.RuntimeEnvOverrides(cfg.Env, cfg.MCPs)
	shareDir, err := resolveKimiShareDir(env, workDir)
	if err != nil {
		return kimiPreparedLaunch{}, err
	}
	prepared := kimiPreparedLaunch{
		Args:     []string{"--wire"},
		WorkDir:  workDir,
		ShareDir: shareDir,
		Cleanup:  func() {},
	}
	success := false
	defer func() {
		if !success {
			prepared.Cleanup()
		}
	}()
	if cfg.DisableNativeTools {
		shareDir, cleanup, isolateErr := prepareKimiIsolatedShareDir(shareDir, workDir)
		if isolateErr != nil {
			return kimiPreparedLaunch{}, isolateErr
		}
		prepared.ShareDir = shareDir
		prepared.addCleanup(cleanup)
	}
	if cfg.Provider != nil {
		if err := prepareKimiProviderConfig(&prepared, cfg); err != nil {
			return kimiPreparedLaunch{}, err
		}
	}
	if env == nil {
		env = make(map[string]string)
	}
	env["KIMI_SHARE_DIR"] = prepared.ShareDir
	permissionMode, err := cfg.EffectivePermissionMode()
	if err != nil {
		return kimiPreparedLaunch{}, err
	}
	switch permissionMode {
	case types.AgentPermissionModeFullAccess:
		prepared.Args = append(prepared.Args, "--yolo")
	}
	prepared.Args = append(prepared.Args,
		"--session", sessionID,
		"--work-dir", workDir,
	)
	if cfg.DisableNativeTools {
		skillsDir := filepath.Join(prepared.ShareDir, "agentteam-skills-empty")
		if err := ensureKimiPrivateDirectory(skillsDir); err != nil {
			return kimiPreparedLaunch{}, fmt.Errorf("kimi: create isolated skills directory: %w", err)
		}
		prepared.Args = append(prepared.Args, "--skills-dir", skillsDir)
		if err := prepareKimiToolRestrictedAgent(&prepared); err != nil {
			return kimiPreparedLaunch{}, err
		}
	}
	if cfg.DisableNativeTools || len(cfg.MCPs) > 0 {
		if err := prepareKimiMCPConfig(&prepared, cfg.MCPs); err != nil {
			return kimiPreparedLaunch{}, err
		}
	}
	prepared.Env = buildRuntimeEnvWithProvider(env, cfg.Provider)
	success = true
	return prepared, nil
}

func prepareKimiToolRestrictedAgent(prepared *kimiPreparedLaunch) error {
	root, err := ensureKimiLaunchConfigDir(prepared)
	if err != nil {
		return err
	}
	prepared.AgentFile = filepath.Join(root, "agent.yaml")
	prepared.SystemPromptFile = filepath.Join(root, "system.md")
	systemPrompt := promptcompiler.KimiRestrictedSystemPrompt()
	agentYAML := "version: 1\nagent:\n  name: agentteam-tool-restricted\n  system_prompt_path: ./system.md\n  tools: []\n  allowed_tools: []\n  exclude_tools: []\n  subagents: {}\n"
	for path, content := range map[string][]byte{
		prepared.SystemPromptFile: []byte(systemPrompt),
		prepared.AgentFile:        []byte(agentYAML),
	} {
		if err := os.WriteFile(path, content, 0o600); err != nil {
			return fmt.Errorf("kimi: write %s: %w", filepath.Base(path), err)
		}
	}
	prepared.Args = append(prepared.Args,
		"--agent-file", prepared.AgentFile,
	)
	return nil
}

func prepareKimiMCPConfig(prepared *kimiPreparedLaunch, mcps []types.NamedMCPConfig) error {
	root, err := ensureKimiLaunchConfigDir(prepared)
	if err != nil {
		return err
	}
	payload, err := kimiMCPConfig(mcps)
	if err != nil {
		return err
	}
	prepared.MCPConfigFile = filepath.Join(root, "mcp.json")
	if err := os.WriteFile(prepared.MCPConfigFile, payload, 0o600); err != nil {
		return fmt.Errorf("kimi: write MCP config: %w", err)
	}
	prepared.Args = append(prepared.Args, "--mcp-config-file", prepared.MCPConfigFile)
	return nil
}

func ensureKimiLaunchConfigDir(prepared *kimiPreparedLaunch) (string, error) {
	switch {
	case prepared == nil:
		return "", fmt.Errorf("kimi: launch config is nil")
	case prepared.AgentFile != "":
		return filepath.Dir(prepared.AgentFile), nil
	case prepared.MCPConfigFile != "":
		return filepath.Dir(prepared.MCPConfigFile), nil
	}
	parent := filepath.Join(prepared.ShareDir, "agentteam-launch")
	if err := os.MkdirAll(parent, 0o700); err != nil {
		return "", fmt.Errorf("kimi: create launch config parent: %w", err)
	}
	if info, err := os.Lstat(parent); err != nil || info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		if err == nil {
			err = fmt.Errorf("path is not a safe directory")
		}
		return "", fmt.Errorf("kimi: inspect launch config parent: %w", err)
	}
	root, err := os.MkdirTemp(parent, "session-*")
	if err != nil {
		return "", fmt.Errorf("kimi: create launch config directory: %w", err)
	}
	if err := os.Chmod(root, 0o700); err != nil {
		_ = os.RemoveAll(root)
		return "", fmt.Errorf("kimi: secure launch config directory: %w", err)
	}
	prepared.addCleanup(func() {
		if err := os.RemoveAll(root); err != nil {
			slog.Warn("kimi: remove launch config directory", "path", root, "error", err)
		}
	})
	return root, nil
}

func resolveKimiShareDir(env map[string]string, workDir string) (string, error) {
	shareDir := strings.TrimSpace(env["KIMI_SHARE_DIR"])
	if shareDir == "" {
		if shellShareDir, ok := runtimeconfig.RuntimeEnvValue("KIMI_SHARE_DIR"); ok {
			shareDir = strings.TrimSpace(shellShareDir)
		}
	}
	if shareDir == "" {
		home, err := os.UserHomeDir()
		if err != nil {
			return "", fmt.Errorf("kimi: resolve home: %w", err)
		}
		shareDir = filepath.Join(home, ".kimi")
	}
	if !filepath.IsAbs(shareDir) {
		shareDir = filepath.Join(workDir, shareDir)
	}
	shareDir = filepath.Clean(shareDir)
	if err := os.MkdirAll(shareDir, 0o700); err != nil {
		return "", fmt.Errorf("kimi: create share directory: %w", err)
	}
	info, err := os.Lstat(shareDir)
	if err != nil {
		return "", fmt.Errorf("kimi: inspect share directory: %w", err)
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		return "", fmt.Errorf("kimi: share directory is not a safe directory")
	}
	return shareDir, nil
}

func validateKimiDirectoryChain(root, target string) error {
	root = filepath.Clean(root)
	target = filepath.Clean(target)
	relative, err := filepath.Rel(root, target)
	if err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
		return fmt.Errorf("path escapes Kimi share directory")
	}
	if info, statErr := os.Lstat(root); statErr != nil || info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		if statErr == nil {
			statErr = fmt.Errorf("root is not a safe directory")
		}
		return statErr
	}
	current := root
	if relative == "." {
		return nil
	}
	for _, component := range strings.Split(relative, string(filepath.Separator)) {
		current = filepath.Join(current, component)
		info, statErr := os.Lstat(current)
		if statErr != nil {
			return statErr
		}
		if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
			return fmt.Errorf("%s is not a safe directory", current)
		}
	}
	resolvedRoot, err := filepath.EvalSymlinks(root)
	if err != nil {
		return err
	}
	resolvedTarget, err := filepath.EvalSymlinks(target)
	if err != nil {
		return err
	}
	resolvedRelative, err := filepath.Rel(resolvedRoot, resolvedTarget)
	if err != nil || resolvedRelative == ".." || strings.HasPrefix(resolvedRelative, ".."+string(filepath.Separator)) {
		return fmt.Errorf("resolved path escapes Kimi share directory")
	}
	return nil
}

func prepareKimiIsolatedShareDir(sourceShareDir, workDir string) (string, func(), error) {
	isolatedShareDir := filepath.Join(filepath.Dir(workDir), "kimi-share")
	cleanup := func() {
		if err := sanitizeKimiIsolatedShareDir(isolatedShareDir); err != nil {
			slog.Warn("kimi: sanitize isolated share directory", "path", isolatedShareDir, "error", err)
		}
	}
	if err := ensureKimiPrivateDirectory(isolatedShareDir); err != nil {
		cleanup()
		return "", nil, fmt.Errorf("kimi: create isolated share directory: %w", err)
	}
	if err := sanitizeKimiIsolatedShareDir(isolatedShareDir); err != nil {
		return "", nil, err
	}
	if err := copyKimiIsolatedConfigIfPresent(
		filepath.Join(sourceShareDir, "config.toml"),
		filepath.Join(isolatedShareDir, "config.toml"),
	); err != nil {
		cleanup()
		return "", nil, err
	}
	if err := copyKimiPrivateDirectoryIfPresent(
		filepath.Join(sourceShareDir, "credentials"),
		filepath.Join(isolatedShareDir, "credentials"),
	); err != nil {
		cleanup()
		return "", nil, err
	}
	return isolatedShareDir, cleanup, nil
}

func copyKimiIsolatedConfigIfPresent(source, target string) error {
	info, err := os.Lstat(source)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("kimi: inspect config.toml: %w", err)
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return fmt.Errorf("kimi: %s is not a safe file", source)
	}
	data, err := os.ReadFile(source)
	if err != nil {
		return fmt.Errorf("kimi: read config.toml: %w", err)
	}
	var sourceConfig map[string]any
	if err := toml.Unmarshal(data, &sourceConfig); err != nil {
		return fmt.Errorf("kimi: decode config.toml: %w", err)
	}
	config := make(map[string]any)
	for _, key := range []string{
		"default_thinking",
		"agentteam_control",
		"telemetry",
	} {
		if value, ok := sourceConfig[key]; ok {
			config[key] = value
		}
	}
	defaultModel, _ := sourceConfig["default_model"].(string)
	defaultModel = strings.TrimSpace(defaultModel)
	models, _ := sourceConfig["models"].(map[string]any)
	model, hasModel := models[defaultModel]
	modelMap, _ := model.(map[string]any)
	providerName, _ := modelMap["provider"].(string)
	providers, _ := sourceConfig["providers"].(map[string]any)
	providerConfig, hasProvider := providers[providerName]
	if defaultModel != "" && hasModel && strings.TrimSpace(providerName) != "" && hasProvider {
		config["default_model"] = defaultModel
		config["models"] = map[string]any{defaultModel: model}
		config["providers"] = map[string]any{providerName: providerConfig}
	}
	sanitized, err := toml.Marshal(config)
	if err != nil {
		return fmt.Errorf("kimi: encode isolated config.toml: %w", err)
	}
	if err := os.WriteFile(target, sanitized, 0o600); err != nil {
		return fmt.Errorf("kimi: write isolated config.toml: %w", err)
	}
	return nil
}

func prepareKimiProviderConfig(prepared *kimiPreparedLaunch, cfg drivercontract.LaunchConfig) error {
	if prepared == nil || cfg.Provider == nil {
		return nil
	}
	prov := cfg.Provider
	apiKey := strings.TrimSpace(prov.APIKey)
	if apiKey == "" {
		if envKey := strings.TrimSpace(prov.EnvKey); envKey != "" {
			if value, ok := runtimeconfig.RuntimeEnvValue(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")
	}
	providerType, err := kimiProviderType(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)
	}
	root, err := ensureKimiLaunchConfigDir(prepared)
	if err != nil {
		return err
	}
	path := filepath.Join(root, "config.toml")
	if err := writeKimiProviderConfigFile(path, payload, 0o600); err != nil {
		return fmt.Errorf("kimi: write provider config: %w", err)
	}
	prepared.Args = append(prepared.Args, "--config-file", path, "--model", alias)
	return nil
}

func sanitizeKimiIsolatedShareDir(path string) error {
	entries, err := os.ReadDir(path)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("kimi: read isolated share directory: %w", err)
	}
	for _, entry := range entries {
		entryPath := filepath.Join(path, entry.Name())
		if entry.Name() == "sessions" {
			info, statErr := os.Lstat(entryPath)
			if statErr != nil {
				return fmt.Errorf("kimi: inspect isolated sessions: %w", statErr)
			}
			if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
				return fmt.Errorf("kimi: isolated sessions path is not a safe directory")
			}
			continue
		}
		if err := os.RemoveAll(entryPath); err != nil {
			return fmt.Errorf("kimi: remove isolated %s: %w", entry.Name(), err)
		}
	}
	return nil
}

func ensureKimiPrivateDirectory(path string) error {
	if err := os.MkdirAll(path, 0o700); err != nil {
		return err
	}
	info, err := os.Lstat(path)
	if err != nil {
		return err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		return fmt.Errorf("%s is not a safe directory", path)
	}
	return os.Chmod(path, 0o700)
}

func copyKimiPrivateDirectoryIfPresent(source, target string) error {
	info, err := os.Lstat(source)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("kimi: inspect %s: %w", filepath.Base(source), err)
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		return fmt.Errorf("kimi: %s is not a safe directory", source)
	}
	if err := ensureKimiPrivateDirectory(target); err != nil {
		return fmt.Errorf("kimi: create isolated %s: %w", filepath.Base(target), err)
	}
	entries, err := os.ReadDir(source)
	if err != nil {
		return fmt.Errorf("kimi: read %s: %w", filepath.Base(source), err)
	}
	for _, entry := range entries {
		if entry.Type()&os.ModeSymlink != 0 || entry.IsDir() {
			return fmt.Errorf("kimi: unsupported entry in %s: %s", source, entry.Name())
		}
		if entry.Name() != "oauth.json" {
			continue
		}
		if err := copyKimiPrivateFileIfPresent(filepath.Join(source, entry.Name()), filepath.Join(target, entry.Name())); err != nil {
			return err
		}
	}
	return nil
}

func copyKimiPrivateFileIfPresent(source, target string) error {
	info, err := os.Lstat(source)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("kimi: inspect %s: %w", filepath.Base(source), err)
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return fmt.Errorf("kimi: %s is not a safe file", source)
	}
	data, err := os.ReadFile(source)
	if err != nil {
		return fmt.Errorf("kimi: read %s: %w", filepath.Base(source), err)
	}
	if err := os.WriteFile(target, data, 0o600); err != nil {
		return fmt.Errorf("kimi: write isolated %s: %w", filepath.Base(target), err)
	}
	return nil
}

func kimiMCPConfig(mcps []types.NamedMCPConfig) ([]byte, error) {
	servers := make(map[string]any)
	for _, mcp := range mcps {
		name := strings.TrimSpace(mcp.Name)
		if name == "" {
			continue
		}
		var server map[string]any
		switch mcp.EffectiveTransport() {
		case types.MCPTransportStdio:
			command := strings.TrimSpace(mcp.Command)
			if command == "" {
				continue
			}
			server = map[string]any{"command": command}
			if len(mcp.Args) > 0 {
				server["args"] = append([]string(nil), mcp.Args...)
			}
			if len(mcp.Env) > 0 {
				server["env"] = runtimeconfig.CloneStringMap(mcp.Env)
			}
			if cwd := strings.TrimSpace(mcp.CWD); cwd != "" {
				server["cwd"] = cwd
			}
		default:
			url := strings.TrimSpace(mcp.URL)
			if url == "" {
				continue
			}
			server = map[string]any{"url": url}
			if len(mcp.Headers) > 0 {
				server["headers"] = runtimeconfig.CloneStringMap(mcp.Headers)
			}
		}
		servers[name] = server
	}
	payload, err := json.Marshal(map[string]any{"mcpServers": servers})
	if err != nil {
		return nil, fmt.Errorf("kimi: encode MCP config: %w", err)
	}
	return append(payload, '\n'), nil
}

func normalizeKimiModelName(model string, prov *drivercontract.ProviderConfig) string {
	model = strings.TrimSpace(model)
	if prov == nil {
		return model
	}
	providerID := strings.Trim(strings.TrimSpace(prov.ID), "/")
	for _, separator := range []string{"/", ":"} {
		prefix := providerID + separator
		if providerID != "" && strings.HasPrefix(model, prefix) {
			return strings.TrimSpace(strings.TrimPrefix(model, prefix))
		}
	}
	return model
}

func kimiProviderType(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_legacy", nil
		case provider.WireAPIResponses:
			return "openai_responses", nil
		}
	}
	haystack := strings.ToLower(prov.ID + " " + prov.Name + " " + prov.BaseURL)
	if strings.Contains(haystack, "anthropic") || strings.Contains(haystack, "claude") {
		return "anthropic", nil
	}
	if strings.Contains(haystack, "responses") {
		return "openai_responses", nil
	}
	if strings.Contains(haystack, "openai") {
		return "openai_legacy", nil
	}
	if strings.Contains(haystack, "kimi") || strings.Contains(haystack, "moonshot") {
		return "kimi", nil
	}
	return "", fmt.Errorf("kimi: unsupported provider protocol for %q", strings.TrimSpace(prov.Name))
}

func kimiInitialPrompt(cfg drivercontract.LaunchConfig, _ bool) string {
	if cfg.CompiledPrompt == nil || !cfg.CompiledPrompt.HasLaunchCarrier() {
		return ""
	}
	return cfg.CompiledPrompt.LaunchSurface("kimiInitialPrompt")
}

func kimiWireUserInput(input drivercontract.RuntimeInput) (any, error) {
	text := input.PromptText()
	if inputID := strings.TrimSpace(input.ID); inputID != "" {
		if text != "" {
			text += "\n\n"
		}
		text += kimiRuntimeInputMarker(inputID)
	}
	if len(input.Attachments) == 0 {
		return text, nil
	}
	parts := make([]map[string]any, 0, len(input.Attachments)+1)
	if strings.TrimSpace(text) != "" {
		parts = append(parts, map[string]any{"type": "text", "text": text})
	}
	for _, attachment := range input.Attachments {
		uri := strings.TrimSpace(attachment.URI)
		if uri == "" {
			continue
		}
		mimeType := strings.ToLower(strings.TrimSpace(attachment.MIME))
		switch {
		case strings.HasPrefix(mimeType, "image/"):
			parts = append(parts, map[string]any{"type": "image_url", "image_url": map[string]any{"url": uri}})
		case strings.HasPrefix(mimeType, "audio/"):
			parts = append(parts, map[string]any{"type": "audio_url", "audio_url": map[string]any{"url": uri}})
		case strings.HasPrefix(mimeType, "video/"):
			parts = append(parts, map[string]any{"type": "video_url", "video_url": map[string]any{"url": uri}})
		default:
			parts = append(parts, map[string]any{"type": "text", "text": "Attachment: " + uri})
		}
	}
	return parts, nil
}

func cloneKimiLaunchConfig(cfg drivercontract.LaunchConfig) drivercontract.LaunchConfig {
	cfg.Env = runtimeconfig.CloneStringMap(cfg.Env)
	return cfg
}

// sortedKimiKeys is used by tests and diagnostics when stable materialization
// ordering matters.
func sortedKimiKeys(values map[string]any) []string {
	keys := make([]string, 0, len(values))
	for key := range values {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	return keys
}