package dispatcher

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"

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

const runtimeStartFingerprintVersion = 6

type runtimeStartFingerprintMaterial struct {
	Version  int                           `json:"version"`
	Runtime  runtimeStartRuntimeMaterial   `json:"runtime"`
	Policy   runtimeStartPolicyMaterial    `json:"policy"`
	Launch   runtimeStartLaunchMaterial    `json:"launch"`
	LLM      runtimeStartLLMMaterial       `json:"llm"`
	Provider *runtimeStartProviderMaterial `json:"provider,omitempty"`
}

type runtimeStartRuntimeMaterial struct {
	Type          string `json:"type"`
	EffectiveType string `json:"effectiveType"`
	Binary        string `json:"binary,omitempty"`
}

type runtimeStartPolicyMaterial struct {
	SchemaVersion  int    `json:"schemaVersion"`
	Generation     int64  `json:"generation"`
	Fingerprint    string `json:"fingerprint"`
	Runtime        string `json:"runtime"`
	RuntimeType    string `json:"runtimeType"`
	PermissionMode string `json:"permissionMode"`
}

type runtimeStartLaunchMaterial struct {
	PermissionMode      string `json:"permissionMode,omitempty"`
	RuntimeDebugEnabled bool   `json:"runtimeDebugEnabled,omitempty"`
	Model               string `json:"model,omitempty"`
	SmallModel          string `json:"smallModel,omitempty"`
	ReasoningEffort     string `json:"reasoningEffort,omitempty"`
}

type runtimeStartLLMMaterial struct {
	Mode string `json:"mode,omitempty"`
}

type runtimeStartProviderMaterial struct {
	BaseURL        string            `json:"baseUrl,omitempty"`
	CredentialKind string            `json:"credentialKind,omitempty"`
	EnvKey         string            `json:"envKey,omitempty"`
	APIKeyDigest   string            `json:"apiKeyDigest,omitempty"`
	Headers        map[string]string `json:"headers,omitempty"`
	WireAPIs       []string          `json:"wireApis,omitempty"`
}

func currentRuntimeStartFingerprint(ag agent.Config, rt runtimeconfig.Runtime) (string, error) {
	policy := currentAgentRuntimePolicy(ag, rt)
	permissionMode := strings.TrimSpace(policy.PermissionMode)

	startCfg := drivercontract.LaunchConfig{
		RuntimeType:         rt.Type,
		Binary:              rt.Binary,
		PermissionMode:      permissionMode,
		RuntimeDebugEnabled: ag.RuntimeDebugEnabled,
		Model:               strings.TrimSpace(ag.Model),
		SmallModel:          strings.TrimSpace(ag.SmallModel),
		ReasoningEffort:     strings.TrimSpace(ag.ReasoningEffort),
	}

	llmResolution, err := profile.ResolveAgentLLM(ag, rt.Type)
	if err != nil {
		return "", err
	}
	if startCfg.Model == "" && llmResolution.Model != "" {
		startCfg.Model = llmResolution.Model
	}
	if startCfg.ReasoningEffort == "" && llmResolution.ReasoningEffort != "" {
		startCfg.ReasoningEffort = llmResolution.ReasoningEffort
	}
	if llmResolution.Upstream != nil {
		upstream := *llmResolution.Upstream
		startCfg.Provider = &drivercontract.ProviderConfig{
			BaseURL:        strings.TrimSpace(upstream.BaseURL),
			APIKey:         strings.TrimSpace(upstream.APIKey),
			CredentialKind: strings.TrimSpace(ag.CredentialKind),
			EnvKey:         strings.TrimSpace(upstream.EnvKey),
			Headers:        upstream.Headers,
			WireAPIs:       append([]string(nil), upstream.WireAPIs...),
		}
	}

	return runtimeStartFingerprint(ag, rt, policy, startCfg), nil
}

func runtimeStartFingerprint(
	ag agent.Config,
	rt runtimeconfig.Runtime,
	policy types.RuntimePolicySnapshot,
	startCfg drivercontract.LaunchConfig,
) string {
	policy = normalizeRuntimePolicySnapshot(policy)
	material := runtimeStartFingerprintMaterial{
		Version: runtimeStartFingerprintVersion,
		Runtime: runtimeStartRuntimeMaterial{
			Type:          strings.TrimSpace(string(rt.Type)),
			EffectiveType: strings.TrimSpace(string(runtimeconfig.EffectiveRuntimeType(rt.Type))),
			Binary:        strings.TrimSpace(runtimeconfig.EffectiveRuntimeBinary(rt.Type, rt.Binary)),
		},
		Policy: runtimeStartPolicyMaterial{
			SchemaVersion:  policy.SchemaVersion,
			Generation:     policy.Generation,
			Fingerprint:    strings.TrimSpace(policy.Fingerprint),
			Runtime:        strings.TrimSpace(policy.Runtime),
			RuntimeType:    strings.TrimSpace(policy.RuntimeType),
			PermissionMode: strings.TrimSpace(policy.PermissionMode),
		},
		Launch: runtimeStartLaunchMaterial{
			PermissionMode:      strings.TrimSpace(startCfg.PermissionMode),
			RuntimeDebugEnabled: startCfg.RuntimeDebugEnabled,
			Model:               strings.TrimSpace(startCfg.Model),
			SmallModel:          strings.TrimSpace(startCfg.SmallModel),
			ReasoningEffort:     strings.TrimSpace(startCfg.ReasoningEffort),
		},
		LLM: runtimeStartLLMMaterial{
			Mode: normalizedAgentLLMMode(ag),
		},
	}
	if startCfg.Provider != nil {
		material.Provider = &runtimeStartProviderMaterial{
			BaseURL:        strings.TrimSpace(startCfg.Provider.BaseURL),
			CredentialKind: strings.TrimSpace(startCfg.Provider.CredentialKind),
			EnvKey:         strings.TrimSpace(startCfg.Provider.EnvKey),
			APIKeyDigest:   digestString(startCfg.Provider.APIKey),
			Headers:        digestStringMap(startCfg.Provider.Headers),
			WireAPIs:       trimmedStringSlice(startCfg.Provider.WireAPIs),
		}
	}
	payload, err := json.Marshal(material)
	if err != nil {
		// All fields are JSON-serializable. Keep this branch explicit so a
		// future field addition cannot silently produce an empty fingerprint.
		panic(fmt.Sprintf("marshal runtime start fingerprint: %v", err))
	}
	sum := sha256.Sum256(payload)
	return "rsf_" + hex.EncodeToString(sum[:])
}

func (d *Dispatcher) restartSessionIfRuntimeStartConfigChanged(sess *interactiveSession, ag agent.Config, rt runtimeconfig.Runtime, binding conversationBinding) (bool, error) {
	if sess == nil {
		return false, nil
	}
	if strings.TrimSpace(string(rt.Type)) == "" {
		return false, nil
	}
	ag = agentConfigForRuntimeStartGuard(ag, binding)
	currentFingerprint, err := currentRuntimeStartFingerprint(ag, rt)
	if err != nil {
		return false, fmt.Errorf("compute runtime start fingerprint: %w", err)
	}

	sess.mu.Lock()
	liveFingerprint := strings.TrimSpace(sess.runtimeStartFingerprint)
	liveRuntimeType := strings.TrimSpace(string(sess.runtime.Type))
	liveRuntimeName := strings.TrimSpace(sess.agent.Runtime)
	liveAgentID := strings.TrimSpace(sess.agentID)
	liveServerAgentID := strings.TrimSpace(sess.serverAgentID)
	sess.mu.Unlock()

	if liveFingerprint == currentFingerprint {
		return false, nil
	}

	// Legacy or test-created in-memory sessions may not have the immutable
	// start fingerprint. Only force a restart in that fallback when explicit
	// runtime identity already proves the live process is stale.
	if liveFingerprint == "" {
		currentRuntimeType := strings.TrimSpace(string(rt.Type))
		if liveRuntimeType != "" && currentRuntimeType != "" && liveRuntimeType != currentRuntimeType {
			slog.Info("restarting live agent session: runtime type changed without start fingerprint",
				"agent_id", liveAgentID,
				"server_agent_id", liveServerAgentID,
				"live_runtime_type", liveRuntimeType,
				"current_runtime_type", currentRuntimeType)
			return d.interruptSessionForRuntimeStartReplacement(sess, ag, rt, currentFingerprint), nil
		}
		if liveRuntimeName != "" && strings.TrimSpace(ag.Runtime) != "" && liveRuntimeName != strings.TrimSpace(ag.Runtime) {
			slog.Info("restarting live agent session: runtime changed without start fingerprint",
				"agent_id", liveAgentID,
				"server_agent_id", liveServerAgentID,
				"live_runtime", liveRuntimeName,
				"current_runtime", strings.TrimSpace(ag.Runtime))
			return d.interruptSessionForRuntimeStartReplacement(sess, ag, rt, currentFingerprint), nil
		}
		return false, nil
	}

	slog.Info("restarting live agent session: runtime start configuration changed",
		"agent_id", liveAgentID,
		"server_agent_id", liveServerAgentID,
		"live_runtime_type", liveRuntimeType,
		"current_runtime_type", strings.TrimSpace(string(rt.Type)))
	return d.interruptSessionForRuntimeStartReplacement(sess, ag, rt, currentFingerprint), nil
}

func agentConfigForRuntimeStartGuard(ag agent.Config, binding conversationBinding) agent.Config {
	if model := strings.TrimSpace(binding.serverModel); model != "" {
		ag.Model = model
	}
	if effort := strings.TrimSpace(binding.serverReasoningEffort); effort != "" {
		ag.ReasoningEffort = effort
	}
	return ag
}

func digestStringMap(input map[string]string) map[string]string {
	if len(input) == 0 {
		return nil
	}
	result := make(map[string]string, len(input))
	for key, value := range input {
		trimmedKey := strings.TrimSpace(key)
		if trimmedKey == "" {
			continue
		}
		result[trimmedKey] = digestString(value)
	}
	if len(result) == 0 {
		return nil
	}
	return result
}

func normalizedAgentLLMMode(ag agent.Config) string {
	mode, _, _, _, _ := agent.NormalizeAgentLLMFields(ag.LLMMode, ag.BaseURL, ag.APIKey, ag.Headers, ag.WireAPI)
	return strings.TrimSpace(mode)
}

func digestString(value string) string {
	trimmed := strings.TrimSpace(value)
	if trimmed == "" {
		return ""
	}
	sum := sha256.Sum256([]byte(trimmed))
	return "sha256:" + hex.EncodeToString(sum[:])
}

func trimmedStringSlice(input []string) []string {
	if len(input) == 0 {
		return nil
	}
	result := make([]string, 0, len(input))
	for _, value := range input {
		if trimmed := strings.TrimSpace(value); trimmed != "" {
			result = append(result, trimmed)
		}
	}
	if len(result) == 0 {
		return nil
	}
	return result
}