Rewrote account switching, fix some minor bugs here and there

pkg/debrid/account/manager.go

@@ -0,0 +1,211 @@
+package account
+
+import (
+	"fmt"
+	"slices"
+	"sync/atomic"
+
+	"github.com/puzpuzpuz/xsync/v4"
+	"github.com/rs/zerolog"
+	"github.com/sirrobot01/decypharr/internal/config"
+	"github.com/sirrobot01/decypharr/internal/request"
+	"github.com/sirrobot01/decypharr/internal/utils"
+	"github.com/sirrobot01/decypharr/pkg/debrid/types"
+	"go.uber.org/ratelimit"
+)
+
+type Manager struct {
+	debrid   string
+	current  atomic.Pointer[Account]
+	accounts *xsync.Map[string, *Account]
+}
+
+func NewManager(debridConf config.Debrid, downloadRL ratelimit.Limiter, logger zerolog.Logger) *Manager {
+	m := &Manager{
+		debrid:   debridConf.Name,
+		accounts: xsync.NewMap[string, *Account](),
+	}
+
+	var firstAccount *Account
+	for idx, token := range debridConf.DownloadAPIKeys {
+		if token == "" {
+			continue
+		}
+		headers := map[string]string{
+			"Authorization": fmt.Sprintf("Bearer %s", token),
+		}
+		account := &Account{
+			Debrid: debridConf.Name,
+			Token:  token,
+			Index:  idx,
+			links:  xsync.NewMap[string, types.DownloadLink](),
+			httpClient: request.New(
+				request.WithRateLimiter(downloadRL),
+				request.WithLogger(logger),
+				request.WithHeaders(headers),
+				request.WithMaxRetries(3),
+				request.WithRetryableStatus(429, 447, 502),
+				request.WithProxy(debridConf.Proxy),
+			),
+		}
+		m.accounts.Store(token, account)
+		if firstAccount == nil {
+			firstAccount = account
+		}
+	}
+	m.current.Store(firstAccount)
+	return m
+}
+
+func (m *Manager) Active() []*Account {
+	activeAccounts := make([]*Account, 0)
+	m.accounts.Range(func(key string, acc *Account) bool {
+		if !acc.Disabled.Load() {
+			activeAccounts = append(activeAccounts, acc)
+		}
+		return true
+	})
+
+	slices.SortFunc(activeAccounts, func(i, j *Account) int {
+		return i.Index - j.Index
+	})
+	return activeAccounts
+}
+
+func (m *Manager) All() []*Account {
+	allAccounts := make([]*Account, 0)
+	m.accounts.Range(func(key string, acc *Account) bool {
+		allAccounts = append(allAccounts, acc)
+		return true
+	})
+
+	slices.SortFunc(allAccounts, func(i, j *Account) int {
+		return i.Index - j.Index
+	})
+	return allAccounts
+}
+
+func (m *Manager) Current() *Account {
+	// Fast path - most common case
+	current := m.current.Load()
+	if current != nil && !current.Disabled.Load() {
+		return current
+	}
+
+	// Slow path - find new current account
+	activeAccounts := m.Active()
+	if len(activeAccounts) == 0 {
+		m.current.Store(nil)
+		return nil
+	}
+
+	newCurrent := activeAccounts[0]
+	m.current.Store(newCurrent)
+	return newCurrent
+}
+
+func (m *Manager) setCurrent(account *Account) {
+	m.current.Store(account)
+}
+
+func (m *Manager) Disable(account *Account) {
+	if account == nil {
+		return
+	}
+
+	account.Disabled.Store(true)
+
+	// If we're disabling the current account, it will be replaced
+	// on the next Current() call - no need to proactively update
+	current := m.current.Load()
+	if current != nil && current.Token == account.Token {
+		// Optional: immediately find replacement
+		activeAccounts := m.Active()
+		if len(activeAccounts) > 0 {
+			m.current.Store(activeAccounts[0])
+		} else {
+			m.current.Store(nil)
+		}
+	}
+}
+
+func (m *Manager) Reset() {
+	m.accounts.Range(func(key string, acc *Account) bool {
+		acc.Disabled.Store(false)
+		return true
+	})
+
+	// Set current to first active account
+	activeAccounts := m.Active()
+	if len(activeAccounts) > 0 {
+		m.current.Store(activeAccounts[0])
+	} else {
+		m.current.Store(nil)
+	}
+}
+
+func (m *Manager) Update(account *Account) {
+	if account != nil {
+		m.accounts.Store(account.Token, account)
+	}
+}
+
+func (m *Manager) GetAccount(token string) (*Account, error) {
+	if token == "" {
+		return nil, fmt.Errorf("token cannot be empty")
+	}
+	acc, ok := m.accounts.Load(token)
+	if !ok {
+		return nil, fmt.Errorf("account not found for token")
+	}
+	return acc, nil
+}
+
+func (m *Manager) GetDownloadLink(fileLink string) (types.DownloadLink, error) {
+	current := m.Current()
+	if current == nil {
+		return types.DownloadLink{}, fmt.Errorf("no active account for debrid service %s", m.debrid)
+	}
+	return current.GetDownloadLink(fileLink)
+}
+
+func (m *Manager) GetAccountFromDownloadLink(downloadLink types.DownloadLink) (*Account, error) {
+	if downloadLink.Link == "" {
+		return nil, fmt.Errorf("cannot get account from empty download link")
+	}
+	if downloadLink.Token == "" {
+		return nil, fmt.Errorf("cannot get account from download link without token")
+	}
+	return m.GetAccount(downloadLink.Token)
+}
+
+func (m *Manager) StoreDownloadLink(downloadLink types.DownloadLink) {
+	if downloadLink.Link == "" || downloadLink.Token == "" {
+		return
+	}
+	account, err := m.GetAccount(downloadLink.Token)
+	if err != nil || account == nil {
+		return
+	}
+	account.StoreDownloadLink(downloadLink)
+}
+
+func (m *Manager) Stats() []map[string]any {
+	stats := make([]map[string]any, 0)
+
+	for _, acc := range m.All() {
+		maskedToken := utils.Mask(acc.Token)
+		accountDetail := map[string]any{
+			"in_use":       acc.Equals(m.Current()),
+			"order":        acc.Index,
+			"disabled":     acc.Disabled.Load(),
+			"token_masked": maskedToken,
+			"username":     acc.Username,
+			"traffic_used": acc.TrafficUsed.Load(),
+			"links_count":  acc.DownloadLinksCount(),
+			"debrid":       acc.Debrid,
+		}
+		stats = append(stats, accountDetail)
+	}
+	return stats
+}