decypharr/pkg/debrid/account/manager.go

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"
)

const (
	MaxDisableCount = 3
)

type Manager struct {
	debrid   string
	current  atomic.Pointer[Account]
	accounts *xsync.Map[string, *Account]
	logger   zerolog.Logger
}

func NewManager(debridConf config.Debrid, downloadRL ratelimit.Limiter, logger zerolog.Logger) *Manager {
	m := &Manager{
		debrid:   debridConf.Name,
		accounts: xsync.NewMap[string, *Account](),
		logger:   logger,
	}

	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 {
		// No active accounts left, try to use disabled ones
		m.logger.Warn().Str("debrid", m.debrid).Msg("No active accounts available, all accounts are disabled")
		allAccounts := m.All()
		if len(allAccounts) == 0 {
			m.logger.Error().Str("debrid", m.debrid).Msg("No accounts configured")
			m.current.Store(nil)
			return nil
		}
		m.current.Store(allAccounts[0])
		return allAccounts[0]
	}

	newCurrent := activeAccounts[0]
	m.current.Store(newCurrent)
	return newCurrent
}

func (m *Manager) Disable(account *Account) {
	if account == nil {
		return
	}

	account.MarkDisabled()

	// 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.Reset()
		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) 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
}

func (m *Manager) CheckAndResetBandwidth() {
	found := false
	m.accounts.Range(func(key string, acc *Account) bool {
		if acc.Disabled.Load() && acc.DisableCount.Load() < MaxDisableCount {
			if err := acc.CheckBandwidth(); err == nil {
				acc.Disabled.Store(false)
				found = true
				m.logger.Info().Str("debrid", m.debrid).Str("token", utils.Mask(acc.Token)).Msg("Re-activated disabled account")
			} else {
				m.logger.Debug().Err(err).Str("debrid", m.debrid).Str("token", utils.Mask(acc.Token)).Msg("Account still disabled")
			}
		}
		return true
	})
	if found {
		// If we re-activated any account, reset current to first active
		activeAccounts := m.Active()
		if len(activeAccounts) > 0 {
			m.current.Store(activeAccounts[0])
		}

	}
}