sub2api/backend/internal/service/windsurf_probe_service.go

package service

import (
	"context"
	"fmt"
	"log/slog"
	"time"

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

type WindsurfProbeService struct {
	cfg         config.WindsurfConfig
	accountRepo AccountRepository
	proxyRepo   ProxyRepository
}

func NewWindsurfProbeService(
	cfg config.WindsurfConfig,
	accountRepo AccountRepository,
	proxyRepo ProxyRepository,
) *WindsurfProbeService {
	return &WindsurfProbeService{
		cfg:         cfg,
		accountRepo: accountRepo,
		proxyRepo:   proxyRepo,
	}
}

type WindsurfProbeResult struct {
	AccountID int64
	Tier      string
	Profile   WindsurfProfileSnapshot
	Status    WindsurfUserStatusSnapshot
	Error     string
}

func (s *WindsurfProbeService) ProbeAccount(ctx context.Context, accountID int64) (*WindsurfProbeResult, error) {
	account, err := s.accountRepo.GetByID(ctx, accountID)
	if err != nil {
		return nil, fmt.Errorf("get account: %w", err)
	}
	if account.Platform != domain.PlatformWindsurf {
		return nil, fmt.Errorf("account %d is not a windsurf account", accountID)
	}

	creds := LoadWindsurfCredentials(account.Credentials)
	if creds.APIKey == "" {
		return nil, fmt.Errorf("account %d has no api_key", accountID)
	}

	proxyURL := ""
	if account.ProxyID != nil {
		proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID)
		if err == nil {
			proxyURL = proxy.URL()
		}
	}

	baseURL := s.cfg.UserStatusBaseURL
	if baseURL == "" {
		baseURL = "https://server.codeium.com"
	}
	client, err := windsurf.NewClient(baseURL, proxyURL)
	if err != nil {
		return nil, fmt.Errorf("create client: %w", err)
	}

	userStatus, err := client.GetUserStatus(ctx, creds.APIKey)
	if err != nil {
		extra := LoadWindsurfExtra(account.Extra)
		extra.Probe.LastProbeAt = time.Now().Format(time.RFC3339)
		extra.Probe.LastProbeError = err.Error()
		account.Extra = StoreWindsurfExtra(extra)
		_ = s.accountRepo.Update(ctx, account)
		return &WindsurfProbeResult{
			AccountID: accountID,
			Error:     err.Error(),
		}, nil
	}

	extra := LoadWindsurfExtra(account.Extra)
	extra.Profile.UserID = userStatus.UserID
	extra.Profile.TeamID = userStatus.TeamID
	extra.Profile.Email = userStatus.Email
	extra.Profile.DisplayName = userStatus.Name
	extra.Profile.PlanName = userStatus.PlanName
	extra.Profile.TierSource = "probe"
	extra.Probe.LastProbeAt = time.Now().Format(time.RFC3339)
	extra.Probe.LastProbeError = ""

	extra.Quota.LastCheckedAt = time.Now().Format(time.RFC3339)
	extra.Quota.LastError = ""
	extra.Quota.DailyPercent = userStatus.DailyPercent
	extra.Quota.WeeklyPercent = userStatus.WeeklyPercent
	extra.Quota.PromptLimit = userStatus.MonthlyPromptCredits
	extra.Quota.PromptUsed = userStatus.UsedPromptCredits
	extra.Quota.FlexLimit = userStatus.MonthlyFlexCredits
	extra.Quota.FlexUsed = userStatus.UsedFlexCredits

	if userStatus.MonthlyPromptCredits != nil && *userStatus.MonthlyPromptCredits > 0 {
		used := float64(0)
		if userStatus.UsedPromptCredits != nil {
			used = *userStatus.UsedPromptCredits
		}
		pct := (used / *userStatus.MonthlyPromptCredits) * 100
		extra.UserStatus.MonthlyPromptCredits = int64(*userStatus.MonthlyPromptCredits)
		extra.UserStatus.UserUsedPromptCredits = int64(used)
		if extra.Quota.DailyPercent == nil {
			extra.Quota.DailyPercent = &pct
		}
	}
	extra.UserStatus.LastFetchedAt = time.Now().Format(time.RFC3339)

	account.Extra = StoreWindsurfExtra(extra)
	if err := s.accountRepo.Update(ctx, account); err != nil {
		slog.Warn("windsurf_probe_save_failed", "account_id", accountID, "error", err)
	}

	return &WindsurfProbeResult{
		AccountID: accountID,
		Tier:      creds.Tier,
		Profile:   extra.Profile,
		Status:    extra.UserStatus,
	}, nil
}

func (s *WindsurfProbeService) ProbeModelCatalog(ctx context.Context, accountID int64) ([]windsurf.ModelInfo, error) {
	account, err := s.accountRepo.GetByID(ctx, accountID)
	if err != nil {
		return nil, fmt.Errorf("get account: %w", err)
	}

	creds := LoadWindsurfCredentials(account.Credentials)
	if creds.APIKey == "" {
		return nil, fmt.Errorf("account %d has no api_key", accountID)
	}

	proxyURL := ""
	if account.ProxyID != nil {
		proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID)
		if err == nil {
			proxyURL = proxy.URL()
		}
	}

	baseURL := s.cfg.UserStatusBaseURL
	if baseURL == "" {
		baseURL = "https://server.codeium.com"
	}
	client, err := windsurf.NewClient(baseURL, proxyURL)
	if err != nil {
		return nil, fmt.Errorf("create client: %w", err)
	}

	models, err := client.ListModels(ctx, creds.APIKey)
	if err != nil {
		return nil, fmt.Errorf("list models: %w", err)
	}

	extra := LoadWindsurfExtra(account.Extra)
	extra.Probe.ModelCatalogAt = time.Now().Format(time.RFC3339)
	account.Extra = StoreWindsurfExtra(extra)
	_ = s.accountRepo.Update(ctx, account)

	return models, nil
}

func (s *WindsurfProbeService) GetRuntime(ctx context.Context, accountID int64) (*WindsurfRuntimeInfo, error) {
	account, err := s.accountRepo.GetByID(ctx, accountID)
	if err != nil {
		return nil, fmt.Errorf("get account: %w", err)
	}
	if account.Platform != domain.PlatformWindsurf {
		return nil, fmt.Errorf("account %d is not a windsurf account", accountID)
	}

	creds := LoadWindsurfCredentials(account.Credentials)
	extra := LoadWindsurfExtra(account.Extra)

	info := &WindsurfRuntimeInfo{
		AccountID:    accountID,
		Tier:         creds.Tier,
		Capabilities: extra.Capabilities,
		ModelMatrix:  extra.ModelMatrix,
	}

	if extra.Probe.LastProbeAt != "" {
		info.LastProbeAt = &extra.Probe.LastProbeAt
	}
	if extra.Refresh.LastStatusRefreshAt != "" {
		info.LastStatusRefreshAt = &extra.Refresh.LastStatusRefreshAt
	}
	if extra.Quota.DailyPercent != nil {
		info.UsagePercent = extra.Quota.DailyPercent
	}
	if extra.UserStatus.MonthlyPromptCredits > 0 {
		info.MonthlyCredits = extra.UserStatus.MonthlyPromptCredits
		info.UsedCredits = extra.UserStatus.UserUsedPromptCredits
	}

	return info, nil
}

type WindsurfRuntimeInfo struct {
	AccountID           int64                              `json:"account_id"`
	Tier                string                             `json:"tier"`
	UsagePercent        *float64                           `json:"usage_percent,omitempty"`
	MonthlyCredits      int64                              `json:"monthly_credits,omitempty"`
	UsedCredits         int64                              `json:"used_credits,omitempty"`
	Capabilities        map[string]WindsurfModelCapability `json:"capabilities,omitempty"`
	ModelMatrix         map[string]WindsurfModelAvail      `json:"model_matrix,omitempty"`
	LastProbeAt         *string                            `json:"last_probe_at,omitempty"`
	LastStatusRefreshAt *string                            `json:"last_status_refresh_at,omitempty"`
}