package dispatcher

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

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

const runtimePolicySchemaVersion = 1

type sessionResumeGuardAction string

const (
	sessionResumeGuardReuse              sessionResumeGuardAction = "reuse"
	sessionResumeGuardClearAndFreshStart sessionResumeGuardAction = "clear_and_fresh_start"
)

type sessionResumeGuardDecision struct {
	Action               sessionResumeGuardAction
	Reason               string
	PersistCurrentPolicy bool
}

func agentRuntimePolicyFromWire(policy *types.RuntimePolicySnapshot) agent.RuntimePolicySnapshot {
	if policy == nil {
		return agent.RuntimePolicySnapshot{}
	}
	return agent.RuntimePolicySnapshot{
		SchemaVersion:  policy.SchemaVersion,
		Generation:     policy.Generation,
		Fingerprint:    strings.TrimSpace(policy.Fingerprint),
		Runtime:        strings.TrimSpace(policy.Runtime),
		RuntimeType:    strings.TrimSpace(policy.RuntimeType),
		PermissionMode: types.NormalizeAgentPermissionMode(policy.PermissionMode),
	}
}

func runtimePolicySnapshotIsZero(policy agent.RuntimePolicySnapshot) bool {
	return policy.SchemaVersion == 0 &&
		policy.Generation == 0 &&
		strings.TrimSpace(policy.Fingerprint) == "" &&
		strings.TrimSpace(policy.Runtime) == "" &&
		strings.TrimSpace(policy.RuntimeType) == "" &&
		strings.TrimSpace(policy.PermissionMode) == ""
}

func currentAgentRuntimePolicy(ag agent.Config, rt runtimeconfig.Runtime) types.RuntimePolicySnapshot {
	permissionMode := currentRuntimePermissionMode(ag.PermissionMode, rt.Type)
	policy := runtimePolicyFromAgentConfig(ag.RuntimePolicy)
	policy = normalizeRuntimePolicySnapshot(policy)

	currentRuntimeType := strings.TrimSpace(string(rt.Type))
	currentRuntime := strings.TrimSpace(policy.Runtime)
	if currentRuntime == "" {
		currentRuntime = strings.TrimSpace(ag.Runtime)
	}
	if currentRuntime == "" {
		currentRuntime = strings.TrimSpace(rt.Name)
	}
	if currentRuntime == "" {
		currentRuntime = currentRuntimeType
	}

	if runtimePolicyMatchesCurrent(policy, permissionMode, currentRuntimeType) {
		if strings.TrimSpace(policy.Runtime) == "" {
			policy.Runtime = currentRuntime
		}
		if strings.TrimSpace(policy.RuntimeType) == "" {
			policy.RuntimeType = currentRuntimeType
		}
		policy.PermissionMode = permissionMode
		return normalizeRuntimePolicySnapshot(policy)
	}

	generation := policy.Generation
	if generation <= 0 {
		generation = 1
	}
	return newRuntimePolicySnapshot(currentRuntime, currentRuntimeType, permissionMode, generation)
}

func currentRuntimePermissionMode(permissionMode string, rt runtimeconfig.RuntimeType) string {
	_ = rt
	return types.NormalizeAgentPermissionMode(permissionMode)
}

func runtimePolicyFromAgentConfig(policy agent.RuntimePolicySnapshot) types.RuntimePolicySnapshot {
	return types.RuntimePolicySnapshot{
		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),
	}
}

func runtimePolicyMatchesCurrent(policy types.RuntimePolicySnapshot, permissionMode string, runtimeType string) bool {
	if runtimePolicySnapshotIsZeroTypes(policy) {
		return false
	}
	if types.NormalizeAgentPermissionMode(policy.PermissionMode) != types.NormalizeAgentPermissionMode(permissionMode) {
		return false
	}
	policyRuntimeType := strings.TrimSpace(policy.RuntimeType)
	return policyRuntimeType == "" || policyRuntimeType == strings.TrimSpace(runtimeType)
}

func newRuntimePolicySnapshot(runtimeName string, runtimeType string, permissionMode string, generation int64) types.RuntimePolicySnapshot {
	if generation <= 0 {
		generation = 1
	}
	policy := types.RuntimePolicySnapshot{
		SchemaVersion:  runtimePolicySchemaVersion,
		Generation:     generation,
		Runtime:        strings.TrimSpace(runtimeName),
		RuntimeType:    strings.TrimSpace(runtimeType),
		PermissionMode: types.NormalizeAgentPermissionMode(permissionMode),
	}
	policy.Fingerprint = runtimePolicyFingerprint(policy.SchemaVersion, policy.Generation, policy.Runtime, policy.RuntimeType, policy.PermissionMode)
	return policy
}

func normalizeRuntimePolicySnapshot(policy types.RuntimePolicySnapshot) types.RuntimePolicySnapshot {
	policy.Fingerprint = strings.TrimSpace(policy.Fingerprint)
	policy.Runtime = strings.TrimSpace(policy.Runtime)
	policy.RuntimeType = strings.TrimSpace(policy.RuntimeType)
	originalPermissionMode := strings.TrimSpace(policy.PermissionMode)
	policy.PermissionMode = types.NormalizeAgentPermissionMode(policy.PermissionMode)
	if policy.SchemaVersion <= 0 {
		policy.SchemaVersion = runtimePolicySchemaVersion
	}
	if policy.Generation <= 0 {
		policy.Generation = 1
	}
	if policy.RuntimeType == "" {
		policy.RuntimeType = policy.Runtime
	}
	if (policy.Fingerprint == "" && (policy.Runtime != "" || policy.RuntimeType != "" || policy.PermissionMode != "")) || policy.PermissionMode != originalPermissionMode {
		policy.Fingerprint = runtimePolicyFingerprint(policy.SchemaVersion, policy.Generation, policy.Runtime, policy.RuntimeType, policy.PermissionMode)
	}
	return policy
}

func runtimePolicyFingerprint(schemaVersion int, generation int64, runtimeName string, runtimeType string, permissionMode string) string {
	if schemaVersion <= 0 {
		schemaVersion = runtimePolicySchemaVersion
	}
	if generation <= 0 {
		generation = 1
	}
	payload := fmt.Sprintf("runtime-policy:v%d:g%d:runtime=%s:type=%s:permission=%s",
		schemaVersion,
		generation,
		strings.TrimSpace(runtimeName),
		strings.TrimSpace(runtimeType),
		strings.TrimSpace(permissionMode),
	)
	sum := sha256.Sum256([]byte(payload))
	return "rp_" + hex.EncodeToString(sum[:])
}

func evaluateSessionResumeGuard(current types.RuntimePolicySnapshot, state sessionResumeState) sessionResumeGuardDecision {
	current = normalizeRuntimePolicySnapshot(current)
	saved, hasSavedPolicy := savedRuntimePolicySnapshot(state)
	if !hasSavedPolicy {
		if strings.TrimSpace(current.PermissionMode) != types.AgentPermissionModeFullAccess {
			return sessionResumeGuardDecision{Action: sessionResumeGuardClearAndFreshStart, Reason: "legacy resume state lacks runtime policy for non-full-access current policy"}
		}
		return sessionResumeGuardDecision{Action: sessionResumeGuardReuse, Reason: "legacy full-access resume state", PersistCurrentPolicy: true}
	}

	saved = normalizeRuntimePolicySnapshot(saved)
	if strings.TrimSpace(saved.Fingerprint) == strings.TrimSpace(current.Fingerprint) && saved.Generation == current.Generation {
		return sessionResumeGuardDecision{Action: sessionResumeGuardReuse, Reason: "runtime policy matched"}
	}
	if permissionModeTightened(saved.PermissionMode, current.PermissionMode) {
		return sessionResumeGuardDecision{Action: sessionResumeGuardClearAndFreshStart, Reason: "runtime permission policy tightened"}
	}
	if strings.TrimSpace(saved.RuntimeType) != "" && strings.TrimSpace(current.RuntimeType) != "" && strings.TrimSpace(saved.RuntimeType) != strings.TrimSpace(current.RuntimeType) {
		return sessionResumeGuardDecision{Action: sessionResumeGuardClearAndFreshStart, Reason: "runtime type changed"}
	}
	if strings.TrimSpace(saved.Runtime) != "" && strings.TrimSpace(current.Runtime) != "" && strings.TrimSpace(saved.Runtime) != strings.TrimSpace(current.Runtime) {
		return sessionResumeGuardDecision{Action: sessionResumeGuardClearAndFreshStart, Reason: "runtime changed"}
	}
	if saved.Generation != current.Generation {
		return sessionResumeGuardDecision{Action: sessionResumeGuardClearAndFreshStart, Reason: "runtime policy generation changed"}
	}
	return sessionResumeGuardDecision{Action: sessionResumeGuardClearAndFreshStart, Reason: "runtime policy fingerprint changed"}
}

func savedRuntimePolicySnapshot(state sessionResumeState) (types.RuntimePolicySnapshot, bool) {
	var policy types.RuntimePolicySnapshot
	if state.LastRuntimePolicy != nil {
		policy = *state.LastRuntimePolicy
	}
	if strings.TrimSpace(policy.Fingerprint) == "" {
		policy.Fingerprint = strings.TrimSpace(state.LastRuntimePolicyFingerprint)
	}
	if policy.Generation <= 0 {
		policy.Generation = state.LastRuntimePolicyGeneration
	}
	hasPolicy := strings.TrimSpace(policy.Fingerprint) != "" || policy.Generation > 0 || strings.TrimSpace(policy.Runtime) != "" || strings.TrimSpace(policy.RuntimeType) != "" || strings.TrimSpace(policy.PermissionMode) != ""
	return policy, hasPolicy
}

func applySessionResumeGuardToState(workDir string, state *sessionResumeState, current types.RuntimePolicySnapshot) (sessionResumeGuardDecision, error) {
	if state == nil {
		return sessionResumeGuardDecision{Action: sessionResumeGuardReuse, Reason: "missing resume state"}, nil
	}
	current = normalizeRuntimePolicySnapshot(current)
	var decision sessionResumeGuardDecision
	err := updateSessionResumeState(workDir, func(latest *sessionResumeState) error {
		// Re-evaluate while holding the store lock. The caller's snapshot may
		// predate a concurrent ownership Put, resume-safety callback, identity
		// update, or deferred append and must never be written back wholesale.
		decision = evaluateSessionResumeGuard(current, *latest)
		switch decision.Action {
		case sessionResumeGuardClearAndFreshStart:
			invalidateResumeStateForRuntimePolicy(latest, current)
		case sessionResumeGuardReuse:
			if decision.PersistCurrentPolicy {
				persistRuntimePolicyOnState(latest, current)
			}
		}
		*state = *latest
		return nil
	})
	return decision, err
}

func guardPersistedResumeState(workDir string, current types.RuntimePolicySnapshot) (sessionResumeGuardDecision, error) {
	state, err := loadSessionResumeState(workDir)
	if err != nil {
		return sessionResumeGuardDecision{}, err
	}
	return applySessionResumeGuardToState(workDir, &state, current)
}

func invalidateResumeStateForRuntimePolicy(state *sessionResumeState, current types.RuntimePolicySnapshot) {
	state.LastSessionID = ""
	state.SessionResumeSafe = false
	if state.LastStartConfig != nil {
		state.LastStartConfig.Runtime = ""
		state.LastStartConfig.PermissionMode = ""
		state.LastStartConfig.InstructionsFingerprint = ""
		state.LastStartConfig.ServerURL = ""
		state.LastStartConfig.AuthToken = ""
		state.LastStartConfig.SessionID = ""
	}
	persistRuntimePolicyOnState(state, current)
}

func persistRuntimePolicyOnState(state *sessionResumeState, current types.RuntimePolicySnapshot) {
	current = normalizeRuntimePolicySnapshot(current)
	policyCopy := current
	state.LastRuntimePolicy = &policyCopy
	state.LastRuntimePolicyFingerprint = strings.TrimSpace(current.Fingerprint)
	state.LastRuntimePolicyGeneration = current.Generation
}

func permissionModeTightened(previous string, current string) bool {
	_, _ = previous, current
	return false
}

func runtimePolicySnapshotIsZeroTypes(policy types.RuntimePolicySnapshot) bool {
	return policy.SchemaVersion == 0 &&
		policy.Generation == 0 &&
		strings.TrimSpace(policy.Fingerprint) == "" &&
		strings.TrimSpace(policy.Runtime) == "" &&
		strings.TrimSpace(policy.RuntimeType) == "" &&
		strings.TrimSpace(policy.PermissionMode) == ""
}

func logResumeGuardDecision(agentID string, decision sessionResumeGuardDecision) {
	if decision.Action != sessionResumeGuardClearAndFreshStart {
		return
	}
	slog.Info("clearing saved session resume state: runtime policy no longer matches current agent policy",
		"agent_id", strings.TrimSpace(agentID),
		"reason", strings.TrimSpace(decision.Reason))
}