// Package state provides thread-safe global state management for the copilot service.
// It maintains configuration, tokens, cached data, and rate limiting information.
package state

import (
	"log/slog"
	"math/rand/v2"
	"sync"
)

// =============================================================================
// TokenManager - Manages multiple token states
// =============================================================================

// TokenManager manages multiple GitHub/Copilot token pairs.
// It provides thread-safe access to token states and selects available tokens.
type TokenManager struct {
	// mu protects all fields from concurrent access
	mu sync.RWMutex

	// tokens stores all token states indexed by GitHub token
	tokens map[string]*TokenState

	// accountType specifies the GitHub account type (individual, business, enterprise)
	accountType string
	// rateLimitSeconds is the rate limit interval between requests
	rateLimitSeconds int
	// rateLimitWait determines whether to wait or error when rate limited
	rateLimitWait bool

	// vsCodeVersion stores the cached VS Code version for compatibility
	vsCodeVersion string
	// copilotChatVersion stores the cached Copilot Chat extension version
	copilotChatVersion string
	// showToken controls whether tokens are logged (for debugging)
	showToken bool
	// selectionAlgorithm controls how tokens are selected.
	// Supported values: "random" (uniform random), "quota_weighted" (weighted by remaining quota).
	// Defaults to "quota_weighted" when empty.
	selectionAlgorithm string
}

// Manager is the singleton instance of TokenManager.
var Manager = &TokenManager{
	tokens: make(map[string]*TokenState),
}

// Initialize initializes the TokenManager with the given tokens.
// It replaces the entire token list with the provided tokens.
// This ensures the manager state matches exactly what's in the token file.
func (m *TokenManager) Initialize(gitHubTokens []string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Create a new map to replace the existing tokens
	newTokens := make(map[string]*TokenState)

	// Create token states for all provided tokens
	for _, token := range gitHubTokens {
		if token != "" {
			// Preserve existing state if available (to maintain refresh goroutines, etc.)
			if existingState, exists := m.tokens[token]; exists {
				newTokens[token] = existingState
			} else {
				newTokens[token] = NewTokenState(token)
			}
		}
	}

	// Replace the tokens map entirely
	m.tokens = newTokens
}

// AddToken adds a new GitHub token to the manager with the associated GitHub username.
// If a token from the same GitHub user already exists, it removes the old one and adds the new one.
// Returns the TokenState for the added token.
func (m *TokenManager) AddToken(gitHubToken string, gitHubUser string) *TokenState {
	m.mu.Lock()
	defer m.mu.Unlock()

	// If this exact token already exists, return it
	if state, exists := m.tokens[gitHubToken]; exists {
		// Update username if provided and current username is empty
		if gitHubUser != "" && state.GitHubUser() == "" {
			state.SetGitHubUser(gitHubUser)
		}
		return state
	}

	// If a username is provided, check for and remove any existing tokens from the same user
	if gitHubUser != "" {
		for token, state := range m.tokens {
			if state.GitHubUser() == gitHubUser && token != gitHubToken {
				// Remove the old token from the same user
				delete(m.tokens, token)
				slog.Debug("Removed old token for user", "user", gitHubUser, "old_token", maskTokenForLog(token))
			}
		}
	}

	// Create and add the new token state
	state := NewTokenState(gitHubToken)
	if gitHubUser != "" {
		state.SetGitHubUser(gitHubUser)
	}
	m.tokens[gitHubToken] = state
	return state
}

// maskTokenForLog masks a token for logging purposes.
func maskTokenForLog(token string) string {
	if len(token) <= 8 {
		return token
	}
	return token[:4] + "****" + token[len(token)-4:]
}

// RemoveToken removes a GitHub token from the manager.
func (m *TokenManager) RemoveToken(gitHubToken string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	delete(m.tokens, gitHubToken)
}

// AvailableToken returns a token state that has available premium quota.
// The selection algorithm is determined by the selectionAlgorithm field:
//   - "random": uniform random selection among available tokens
//   - "quota_weighted" (default): weighted random based on premium quota remaining
//
// Only returns enabled tokens with premium quota and a valid Copilot token.
func (m *TokenManager) AvailableToken() *TokenState {
	m.mu.Lock()
	defer m.mu.Unlock()

	var availableStates []*TokenState
	for _, state := range m.tokens {
		if state.Enabled() &&
			state.HasPremiumQuota() &&
			state.CopilotToken() != "" {
			availableStates = append(availableStates, state)
		}
	}

	if len(availableStates) == 0 {
		return nil
	}

	// Read selectionAlgorithm directly (not via getter) since m.mu is already held
	switch m.selectionAlgorithm {
	case "quota_weighted":
		return m.selectByQuotaWeight(availableStates)
	default:
		// Default to simple random selection
		return availableStates[rand.IntN(len(availableStates))]
	}
}

// AvailableTokenForModel returns a token state that has available premium quota
// and whose subscription tier meets the model's minimum tier requirement.
func (m *TokenManager) AvailableTokenForModel(model string) *TokenState {
	m.mu.Lock()
	defer m.mu.Unlock()

	requiredTier := TierForModel(model)

	var availableStates []*TokenState
	for _, state := range m.tokens {
		if !state.Enabled() ||
			state.CopilotToken() == "" ||
			!state.HasPremiumQuota() {
			continue
		}
		if state.PlanTier() < requiredTier {
			continue
		}
		availableStates = append(availableStates, state)
	}

	if len(availableStates) == 0 {
		return nil
	}

	switch m.selectionAlgorithm {
	case "quota_weighted":
		return m.selectByQuotaWeight(availableStates)
	default:
		return availableStates[rand.IntN(len(availableStates))]
	}
}

// selectByQuotaWeight selects a token using weighted random based on premium quota percentage remaining.
// Tokens with higher remaining percentage are proportionally more likely to be selected.
// This distributes load so that tokens with more quota headroom absorb more requests,
// naturally balancing usage across all tokens.
// Caller must hold m.mu.
func (m *TokenManager) selectByQuotaWeight(availableStates []*TokenState) *TokenState {
	var weights []float64
	var totalWeight float64
	for _, state := range availableStates {
		quota := state.PremiumQuota()
		w := quota.PercentRemaining
		if quota.Unlimited {
			// Unlimited tokens always get the maximum weight
			w = 100.0
		} else if w <= 0 {
			// Give a minimal weight so it can still be picked as a fallback
			w = 1.0
		}
		weights = append(weights, w)
		totalWeight += w
	}

	r := rand.Float64() * totalWeight
	var cumulative float64
	var selected *TokenState
	for i, state := range availableStates {
		cumulative += weights[i]
		if r < cumulative {
			selected = state
			break
		}
	}
	// Safety fallback (should not happen due to float precision, but just in case)
	if selected == nil {
		selected = availableStates[len(availableStates)-1]
	}

	return selected
}

// AnyToken returns any enabled token state from the manager.
func (m *TokenManager) AnyToken() *TokenState {
	m.mu.Lock()
	defer m.mu.Unlock()

	for _, state := range m.tokens {
		if state.Enabled() && state.CopilotToken() != "" {
			return state
		}
	}
	return nil
}

// AllTokens returns all token states.
func (m *TokenManager) AllTokens() []*TokenState {
	m.mu.RLock()
	defer m.mu.RUnlock()

	states := make([]*TokenState, 0, len(m.tokens))
	for _, state := range m.tokens {
		states = append(states, state)
	}
	return states
}

// TokenCount returns the number of tokens managed by this manager.
func (m *TokenManager) TokenCount() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.tokens)
}

// AccountType returns the configured GitHub account type thread-safely.
func (m *TokenManager) AccountType() string {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.accountType
}

// SetAccountType updates the GitHub account type thread-safely.
func (m *TokenManager) SetAccountType(accountType string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.accountType = accountType
}

// RateLimitSeconds returns the rate limit interval in seconds thread-safely.
func (m *TokenManager) RateLimitSeconds() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.rateLimitSeconds
}

// SetRateLimitSeconds updates the rate limit interval thread-safely.
func (m *TokenManager) SetRateLimitSeconds(rateLimitSeconds int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.rateLimitSeconds = rateLimitSeconds
}

// RateLimitWait returns whether to wait on rate limit hits thread-safely.
func (m *TokenManager) RateLimitWait() bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.rateLimitWait
}

// SetRateLimitWait updates the rate limit wait behavior thread-safely.
func (m *TokenManager) SetRateLimitWait(rateLimitWait bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.rateLimitWait = rateLimitWait
}

// VSCodeVersion returns the cached VS Code version thread-safely.
func (m *TokenManager) VSCodeVersion() string {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.vsCodeVersion
}

// SetVSCodeVersion updates the cached VS Code version thread-safely.
func (m *TokenManager) SetVSCodeVersion(vsCodeVersion string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.vsCodeVersion = vsCodeVersion
}

// CopilotChatVersion returns the cached Copilot Chat extension version thread-safely.
func (m *TokenManager) CopilotChatVersion() string {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.copilotChatVersion
}

// SetCopilotChatVersion updates the cached Copilot Chat extension version thread-safely.
func (m *TokenManager) SetCopilotChatVersion(copilotChatVersion string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.copilotChatVersion = copilotChatVersion
}

// ShowToken returns whether tokens should be logged thread-safely.
func (m *TokenManager) ShowToken() bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.showToken
}

// SetShowToken updates the token logging preference thread-safely.
func (m *TokenManager) SetShowToken(showToken bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.showToken = showToken
}

// SelectionAlgorithm returns the current token selection algorithm thread-safely.
// Returns "random" when the field is empty (default behavior).
func (m *TokenManager) SelectionAlgorithm() string {
	m.mu.RLock()
	defer m.mu.RUnlock()
	if m.selectionAlgorithm == "" {
		return "random"
	}
	return m.selectionAlgorithm
}

// SetSelectionAlgorithm updates the token selection algorithm thread-safely.
// Supported values: "random", "quota_weighted".
func (m *TokenManager) SetSelectionAlgorithm(algorithm string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.selectionAlgorithm = algorithm
}