sub2api/backend/internal/service/windsurf_tier_access_service.go

package service

import (
	"context"
	"fmt"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/domain"
	"github.com/Wei-Shaw/sub2api/internal/pkg/windsurf"
)

// WindsurfTierAccessRow describes account-pool availability for one model.
//
// Counts are exclusive: free + pro + trial sums to total schedulable accounts
// that can serve the model; blocked excludes accounts whose schedulable=false
// or capability check failed. Total = free + pro + trial.
type WindsurfTierAccessRow struct {
	Model           string `json:"model"`
	Provider        string `json:"provider"`
	EmulationFlavor string `json:"emulation_flavor"`
	Free            int    `json:"free"`
	Pro             int    `json:"pro"`
	Trial           int    `json:"trial"`
	Blocked         int    `json:"blocked"`
	Total           int    `json:"total"`
}

// WindsurfTierAccessSnapshot is the cacheable result of a tier-access scan.
type WindsurfTierAccessSnapshot struct {
	GeneratedAt time.Time               `json:"generated_at"`
	Accounts    int                     `json:"accounts_considered"`
	Rows        []WindsurfTierAccessRow `json:"rows"`
}

// WindsurfTierAccessService aggregates per-model availability from the
// account pool. The Snapshot result is cached for cacheTTL to keep this
// cheap when called from a busy admin dashboard.
type WindsurfTierAccessService struct {
	accountRepo AccountRepository
	cacheTTL    time.Duration

	cache atomic.Pointer[WindsurfTierAccessSnapshot]
	mu    sync.Mutex // guards rebuild
}

// NewWindsurfTierAccessService creates a service with a default 60s cache.
func NewWindsurfTierAccessService(accountRepo AccountRepository) *WindsurfTierAccessService {
	return &WindsurfTierAccessService{
		accountRepo: accountRepo,
		cacheTTL:    60 * time.Second,
	}
}

// Snapshot returns the latest tier-access snapshot, rebuilding from the
// repository when the cache is stale or absent. Concurrent callers during
// a rebuild get the freshly generated snapshot; only one rebuild runs at a
// time.
func (s *WindsurfTierAccessService) Snapshot(ctx context.Context) (*WindsurfTierAccessSnapshot, error) {
	if cached := s.cache.Load(); cached != nil && time.Since(cached.GeneratedAt) < s.cacheTTL {
		return cached, nil
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	// Re-check after acquiring lock — another goroutine may have rebuilt.
	if cached := s.cache.Load(); cached != nil && time.Since(cached.GeneratedAt) < s.cacheTTL {
		return cached, nil
	}

	snap, err := s.build(ctx)
	if err != nil {
		return nil, err
	}
	s.cache.Store(snap)
	return snap, nil
}

func (s *WindsurfTierAccessService) build(ctx context.Context) (*WindsurfTierAccessSnapshot, error) {
	accounts, err := s.accountRepo.ListByPlatform(ctx, domain.PlatformWindsurf)
	if err != nil {
		return nil, fmt.Errorf("list windsurf accounts: %w", err)
	}

	byModel := make(map[string]*tierCounter)
	getCounter := func(model string) *tierCounter {
		c, ok := byModel[model]
		if !ok {
			c = &tierCounter{}
			byModel[model] = c
		}
		return c
	}

	considered := 0
	for i := range accounts {
		acct := &accounts[i]
		creds := LoadWindsurfCredentials(acct.Credentials)
		extra := LoadWindsurfExtra(acct.Extra)
		if creds.APIKey == "" {
			continue // un-registered account; cannot serve traffic
		}
		considered++

		tierBucket := tierBucketFor(creds.Tier)
		schedulable := acct.IsSchedulable()

		// 1) Walk the account's allowedModels (authoritative when present).
		seen := make(map[string]struct{})
		for _, am := range extra.UserStatus.AllowedModels {
			model := am.ModelKey
			if model == "" {
				model = am.Alias
			}
			if model == "" {
				continue
			}
			seen[model] = struct{}{}
			c := getCounter(model)
			if !schedulable {
				c.blocked++
				continue
			}
			capCheck := extra.Capabilities[model]
			if !capabilityOK(capCheck) {
				c.blocked++
				continue
			}
			incTier(c, tierBucket)
		}

		// 2) Fall back to capability map for any model not already counted
		//    (older accounts may have probe data without allowedModels).
		for model, capCheck := range extra.Capabilities {
			if _, ok := seen[model]; ok {
				continue
			}
			c := getCounter(model)
			if !schedulable || !capabilityOK(capCheck) {
				c.blocked++
				continue
			}
			incTier(c, tierBucket)
		}
	}

	rows := make([]WindsurfTierAccessRow, 0, len(byModel))
	for model, c := range byModel {
		meta := windsurf.GetModelInfo(model)
		row := WindsurfTierAccessRow{
			Model:   model,
			Free:    c.free,
			Pro:     c.pro,
			Trial:   c.trial,
			Blocked: c.blocked,
			Total:   c.free + c.pro + c.trial,
		}
		if meta != nil {
			row.Provider = meta.Provider
			row.EmulationFlavor = windsurf.ResolveEmulationFlavor(meta)
		}
		rows = append(rows, row)
	}
	sort.Slice(rows, func(i, j int) bool {
		if rows[i].Total != rows[j].Total {
			return rows[i].Total > rows[j].Total
		}
		return rows[i].Model < rows[j].Model
	})

	return &WindsurfTierAccessSnapshot{
		GeneratedAt: time.Now(),
		Accounts:    considered,
		Rows:        rows,
	}, nil
}

// tierCounter is the per-model tally used during a snapshot build.
type tierCounter struct {
	free, pro, trial, blocked int
}

func capabilityOK(c WindsurfModelCapability) bool {
	if c.Reason == "not_entitled" {
		return false
	}
	return c.Available
}

func tierBucketFor(tier string) string {
	switch tier {
	case "pro":
		return "pro"
	case "trial":
		return "trial"
	case "free":
		return "free"
	default:
		// Unknown tiers (legacy / pre-probe accounts) bucket as free for
		// display purposes — they're typically free until probed.
		return "free"
	}
}

func incTier(c *tierCounter, bucket string) {
	switch bucket {
	case "pro":
		c.pro++
	case "trial":
		c.trial++
	default:
		c.free++
	}
}