package daemon

import (
	"context"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/runtimecontrol"
)

// WarmPoolKey scopes warm state to the same same-node/session-generation
// boundary as link grants. It is an internal map key only and must not be used
// as a metrics label.
type WarmPoolKey struct {
	WorkspaceID       string
	MachineID         string
	SessionGeneration string
	NodeID            string
}

type WarmLease struct {
	Key         WarmPoolKey
	Origin      runtimecontrol.RemoteRequestOrigin
	ProfileName string
	ActivatedAt time.Time
	ExpiresAt   time.Time
}

type warmCooldownState struct {
	Until        time.Time
	FailureCount int
}

type WarmPoolController struct {
	leases    map[warmPoolMapKey]WarmLease
	cooldowns map[warmPoolMapKey]warmCooldownState
}

type warmPoolMapKey struct {
	WarmPoolKey
	ProfileName string
}

type warmPoolStats struct {
	Connecting int
	Idle       int
	Assigned   int
	Total      int
}

func newWarmPoolController() *WarmPoolController {
	return &WarmPoolController{
		leases:    make(map[warmPoolMapKey]WarmLease),
		cooldowns: make(map[warmPoolMapKey]warmCooldownState),
	}
}

// OnRequestActivity refreshes a user-active warm lease and opportunistically
// opens one warm carrier. It intentionally does not wait for Desktop to connect,
// so the first user request keeps its existing LinkPolicy/fallback/no-retry
// behavior; proactive warm failures are converted into cooldown telemetry rather
// than caller-visible errors.
func (manager *ElasticRequestLinkManager) OnRequestActivity(ctx context.Context, client *clientState, descriptor RemoteEndpointDescriptor, request runtimecontrol.RemoteRequestContext) {
	key, profileName, ok := manager.PrepareWarmPoolActivity(ctx, client, descriptor, request)
	if !ok {
		return
	}
	manager.ReconcileWarmPool(client, key, profileName)
}

func (manager *ElasticRequestLinkManager) PrepareWarmPoolActivity(ctx context.Context, client *clientState, descriptor RemoteEndpointDescriptor, request runtimecontrol.RemoteRequestContext) (WarmPoolKey, string, bool) {
	if manager == nil {
		return WarmPoolKey{}, "", false
	}
	if request.Origin == "" && request.Cause == "" {
		request = runtimecontrol.RemoteRequestContextFromContext(ctx)
	} else {
		request = runtimecontrol.RemoteRequestContextFromContext(runtimecontrol.WithRemoteRequestContext(ctx, request))
	}
	if request.Origin != runtimecontrol.RemoteRequestOriginUserInteractive {
		manager.recordWarmEvent(string(request.Origin), defaultWarmEndpointProfileName, "skipped_background", string(request.Origin))
		return WarmPoolKey{}, "", false
	}
	if !manager.warmTransportEligible(client) {
		manager.recordWarmEvent(string(request.Origin), defaultWarmEndpointProfileName, "skipped_ineligible", "transport")
		return WarmPoolKey{}, "", false
	}
	key := manager.warmPoolKeyFromClient(client)
	if !key.valid() {
		manager.recordWarmEvent(string(request.Origin), defaultWarmEndpointProfileName, "skipped_ineligible", "session")
		return WarmPoolKey{}, "", false
	}
	profileName := defaultWarmEndpointProfileName
	activatedAt := manager.config.Now().UTC()
	manager.mu.Lock()
	manager.ensureWarmControllerLocked()
	leaseKey := warmPoolMapKey{WarmPoolKey: key, ProfileName: profileName}
	manager.warm.leases[leaseKey] = WarmLease{
		Key:         key,
		Origin:      runtimecontrol.RemoteRequestOriginUserInteractive,
		ProfileName: profileName,
		ActivatedAt: activatedAt,
		ExpiresAt:   activatedAt.Add(manager.config.WarmActiveTTL),
	}
	manager.recordWarmEvent(string(request.Origin), profileName, "lease_refreshed", request.Cause)
	manager.mu.Unlock()

	_ = descriptor // reserved for later profile selection without changing the public signature.
	return key, profileName, true
}

func (manager *ElasticRequestLinkManager) ReconcileWarmPool(client *clientState, key WarmPoolKey, profileName string) {
	if manager == nil {
		return
	}
	key = key.normalized()
	profileName = strings.TrimSpace(profileName)
	if profileName == "" {
		profileName = defaultWarmEndpointProfileName
	}
	if !key.valid() || !manager.warmTransportEligible(client) {
		manager.recordWarmEvent("user_interactive", profileName, "skipped_ineligible", "transport")
		return
	}
	spec, err := manager.WarmLinkGrantSpec(profileName)
	if err != nil {
		manager.recordWarmEvent("user_interactive", profileName, "skipped_ineligible", "profile")
		return
	}

	var open *ElasticRequestLinkOpenCommand
	var toClose []elasticRequestLinkSender
	var skipped bool
	var skipReason string
	now := manager.config.Now().UTC()
	manager.mu.Lock()
	manager.ensureWarmControllerLocked()
	leaseKey := warmPoolMapKey{WarmPoolKey: key, ProfileName: spec.ProfileName}
	lease := manager.warm.leases[leaseKey]
	activeLease := lease.Active(now)
	targetFloor := 0
	if activeLease && !manager.config.EmergencyDisabled {
		targetFloor = manager.config.WarmTargetIdleFloor
	}
	if targetFloor <= 0 {
		closeSenders, _ := manager.drainSurplusWarmLinksLocked(key, spec.ProfileName, "lease_expired")
		toClose = append(toClose, closeSenders...)
		manager.recordWarmEvent("user_interactive", spec.ProfileName, "drained", "lease_expired")
		manager.recordPoolLocked()
		manager.mu.Unlock()
		for _, sender := range toClose {
			if sender != nil {
				_ = sender.Close()
			}
		}
		return
	}
	stats := manager.warmStatsLocked(key, spec.ProfileName)
	switch {
	case stats.Idle+stats.Connecting >= targetFloor:
		manager.recordWarmEvent("user_interactive", spec.ProfileName, "exists", "ok")
		skipped = true
	case stats.Total >= manager.config.WarmMaxLinks:
		manager.recordWarmEvent("user_interactive", spec.ProfileName, "skipped_capacity", "capacity")
		skipped = true
	case manager.activeCountLocked(key.MachineID) >= manager.config.MaxPerMachine:
		manager.recordWarmEvent("user_interactive", spec.ProfileName, "skipped_capacity", "capacity")
		skipped = true
	case manager.warmCooldownActiveLocked(leaseKey, now):
		manager.recordWarmEvent("user_interactive", spec.ProfileName, "cooldown", "backoff")
		skipped = true
	default:
		openCommand, _, createErr := manager.createGrantForSpecLocked(client, spec)
		if createErr != nil {
			manager.recordWarmEvent("user_interactive", spec.ProfileName, "skipped_capacity", redactedElasticErrorReason(createErr))
			skipped = true
			skipReason = redactedElasticErrorReason(createErr)
		} else {
			command := openCommand
			open = &command
			manager.recordWarmEvent("user_interactive", spec.ProfileName, "grant_created", "ok")
		}
	}
	manager.mu.Unlock()
	for _, sender := range toClose {
		if sender != nil {
			_ = sender.Close()
		}
	}
	if skipped {
		_ = skipReason
		return
	}
	if open == nil {
		return
	}
	if err := manager.sendWarmOpen(client, *open); err != nil {
		manager.CleanupLinkGrant(open.LinkID, "send_open_failed")
		manager.recordWarmEvent("user_interactive", open.ProfileName, "cooldown", "send_open_failed")
	}
}

func (manager *ElasticRequestLinkManager) sendWarmOpen(client *clientState, open ElasticRequestLinkOpenCommand) error {
	if manager.config.WarmOpenSender != nil {
		return manager.config.WarmOpenSender(client, open)
	}
	return client.sessionLinkSender().Send(open)
}

func (manager *ElasticRequestLinkManager) warmTransportEligible(client *clientState) bool {
	if manager == nil || manager.config.EmergencyDisabled || client == nil {
		return false
	}
	if strings.TrimSpace(client.currentSessionGeneration()) == "" {
		return false
	}
	if !client.supportsElasticRequestLinks() {
		return false
	}
	spec, err := manager.WarmLinkGrantSpec(defaultWarmEndpointProfileName)
	if err != nil {
		return false
	}
	return manager.linkGrantSpecSupportedByClient(client, spec)
}

func (manager *ElasticRequestLinkManager) warmPoolKeyFromClient(client *clientState) WarmPoolKey {
	if manager == nil || client == nil {
		return WarmPoolKey{}
	}
	return WarmPoolKey{
		WorkspaceID:       strings.TrimSpace(client.workspaceID),
		MachineID:         strings.TrimSpace(client.machineID),
		SessionGeneration: strings.TrimSpace(client.currentSessionGeneration()),
		NodeID:            strings.TrimSpace(manager.config.NodeID),
	}.normalized()
}

func (key WarmPoolKey) normalized() WarmPoolKey {
	key.WorkspaceID = strings.TrimSpace(key.WorkspaceID)
	key.MachineID = strings.TrimSpace(key.MachineID)
	key.SessionGeneration = strings.TrimSpace(key.SessionGeneration)
	key.NodeID = strings.TrimSpace(key.NodeID)
	return key
}

func (key WarmPoolKey) valid() bool {
	key = key.normalized()
	return key.WorkspaceID != "" && key.MachineID != "" && key.SessionGeneration != "" && key.NodeID != ""
}

func (lease WarmLease) Active(now time.Time) bool {
	if lease.ExpiresAt.IsZero() {
		return false
	}
	if now.IsZero() {
		now = time.Now().UTC()
	}
	return now.Before(lease.ExpiresAt) || now.Equal(lease.ExpiresAt)
}

func (manager *ElasticRequestLinkManager) ensureWarmControllerLocked() {
	if manager.warm == nil {
		manager.warm = newWarmPoolController()
	}
}

func (manager *ElasticRequestLinkManager) warmStatsLocked(key WarmPoolKey, profileName string) warmPoolStats {
	key = key.normalized()
	profileName = strings.TrimSpace(profileName)
	stats := warmPoolStats{}
	for _, grantState := range manager.grants {
		if grantState == nil || grantState.consumed || !warmGrantMatchesKey(grantState.grant, key, profileName) {
			continue
		}
		stats.Connecting++
		stats.Total++
	}
	for _, link := range manager.links {
		if link == nil || !warmGrantMatchesKey(link.grant, key, profileName) {
			continue
		}
		if link.currentAssignment != "" {
			stats.Assigned++
		} else {
			stats.Idle++
		}
		stats.Total++
	}
	return stats
}

func warmGrantMatchesKey(grant ElasticRequestLinkGrant, key WarmPoolKey, profileName string) bool {
	key = key.normalized()
	return grant.Purpose == ElasticLinkPurposeWarm &&
		strings.TrimSpace(grant.ProfileName) == strings.TrimSpace(profileName) &&
		strings.TrimSpace(grant.WorkspaceID) == key.WorkspaceID &&
		strings.TrimSpace(grant.MachineID) == key.MachineID &&
		strings.TrimSpace(grant.SessionGeneration) == key.SessionGeneration &&
		strings.TrimSpace(grant.NodeHint) == key.NodeID
}

func (manager *ElasticRequestLinkManager) drainSurplusWarmLinksLocked(key WarmPoolKey, profileName string, reason string) ([]elasticRequestLinkSender, int) {
	var toClose []elasticRequestLinkSender
	closed := 0
	for linkID, grantState := range manager.grants {
		if grantState == nil || grantState.consumed || !warmGrantMatchesKey(grantState.grant, key, profileName) {
			continue
		}
		delete(manager.grants, linkID)
		manager.recordWarmEvent("user_interactive", profileName, "closed", reason)
	}
	for linkID, link := range manager.links {
		if link == nil || !warmGrantMatchesKey(link.grant, key, profileName) {
			continue
		}
		if link.currentAssignment == "" {
			toClose = append(toClose, link.sender)
			manager.removeLinkLocked(linkID)
			closed++
			manager.recordElasticEvent("cleanup_closed_link", reason)
			manager.recordWarmEvent("user_interactive", profileName, "closed", reason)
			continue
		}
		link.draining = true
		manager.recordWarmEvent("user_interactive", profileName, "drained", reason)
	}
	return toClose, closed
}

func (manager *ElasticRequestLinkManager) warmCooldownActiveLocked(key warmPoolMapKey, now time.Time) bool {
	manager.ensureWarmControllerLocked()
	cooldown := manager.warm.cooldowns[key]
	return !cooldown.Until.IsZero() && now.Before(cooldown.Until)
}

func (manager *ElasticRequestLinkManager) recordWarmFailureLocked(grant ElasticRequestLinkGrant, reason string, now time.Time) {
	if grant.Purpose != ElasticLinkPurposeWarm {
		return
	}
	manager.ensureWarmControllerLocked()
	key := warmPoolMapKey{
		WarmPoolKey: WarmPoolKey{
			WorkspaceID:       grant.WorkspaceID,
			MachineID:         grant.MachineID,
			SessionGeneration: grant.SessionGeneration,
			NodeID:            grant.NodeHint,
		}.normalized(),
		ProfileName: strings.TrimSpace(grant.ProfileName),
	}
	cooldown := manager.warm.cooldowns[key]
	cooldown.FailureCount++
	delay := manager.config.WarmCooldownBase
	// Cooldown applies only to warm carrier setup failures that imply repeated
	// proactive dials would churn. Normal idle/max-age/drain cleanup is not a
	// failure and therefore must not extend backoff or affect the current user
	// request's route/fallback result.
	for i := 1; i < cooldown.FailureCount; i++ {
		delay *= 2
		if delay >= manager.config.WarmCooldownMax {
			delay = manager.config.WarmCooldownMax
			break
		}
	}
	if delay > manager.config.WarmCooldownMax {
		delay = manager.config.WarmCooldownMax
	}
	cooldown.Until = now.Add(delay)
	manager.warm.cooldowns[key] = cooldown
	manager.recordWarmEvent("user_interactive", grant.ProfileName, "cooldown", reason)
}

func warmFailureCounts(reason string) bool {
	switch strings.TrimSpace(reason) {
	case "connect_timeout", "send_open_failed", "token_reject", "dial_failed", "protocol_reject",
		string(ElasticRequestLinkRejectExpired),
		string(ElasticRequestLinkRejectReplayed),
		string(ElasticRequestLinkRejectNodeMismatch),
		string(ElasticRequestLinkRejectRegistryMismatch),
		string(ElasticRequestLinkRejectProtocolMismatch),
		string(ElasticRequestLinkRejectInvalidToken),
		string(ElasticRequestLinkRejectMalformedToken):
		return true
	default:
		return false
	}
}