thread-master/apps/backend/internal/module/usage/usecase/service.go

package usecase

import (
	"context"
	"fmt"
	"strings"
	"time"

	"apps/backend/internal/module/usage/domain"

	"github.com/google/uuid"
)

// KeyResolver decides platform vs byok for a meter.
type KeyResolver interface {
	// Resolve returns key_mode and whether a usable key exists.
	Resolve(ctx context.Context, uid int64, meter string) (keyMode string, hasKey bool, err error)
}

// StaticResolver for tests / simple wiring.
type StaticResolver struct {
	// ByUIDMeter: "uid:meter" -> keyMode (platform|byok); empty means no key
	Map map[string]string
}

func (s *StaticResolver) Resolve(_ context.Context, uid int64, meter string) (string, bool, error) {
	if s == nil || s.Map == nil {
		return "", false, nil
	}
	k := s.Map[fmt.Sprintf("%d:%s", uid, meter)]
	if k == "" {
		return "", false, nil
	}
	return k, true, nil
}

type Service struct {
	Repo     domain.Repository
	Resolver KeyResolver
	// Platform capacity / rate limit. nil = always allow platform.
	// When denied, only BYOK (already preferred in Resolver) can proceed.
	Gate PlatformGate
}

func New(repo domain.Repository, resolver KeyResolver) *Service {
	return &Service{Repo: repo, Resolver: resolver}
}

// RecordCall preserves the audit event after a call. Platform credits were already
// reserved by PrepareCall; this method must not charge them a second time.
func (s *Service) RecordCall(ctx context.Context, uid int64, meter, keyMode, label, source string) (*domain.Event, error) {
	if keyMode != domain.KeyModePlatform && keyMode != domain.KeyModeByok {
		return nil, fmt.Errorf("invalid key_mode")
	}
	credits := 0
	if keyMode == domain.KeyModePlatform {
		credits = domain.MeterCost(meter)
	}
	now := domain.NowNano()
	e := &domain.Event{
		ID: uuid.NewString(), UID: uid, Meter: meter, Credits: credits,
		KeyMode: keyMode, Label: label, Source: source,
		MonthKey: domain.MonthKey(now), CreatedAt: now,
	}
	if err := s.Repo.InsertEvent(ctx, e); err != nil {
		return nil, err
	}
	return e, nil
}

// PrepareCall checks quota / key before external call; returns key_mode.
// Platform path: capacity gate first, then plan credits. BYOK skips both.
func (s *Service) PrepareCall(ctx context.Context, uid int64, meter string) (keyMode string, err error) {
	if s.Resolver == nil {
		return "", domain.ErrNoKey
	}
	mode, has, err := s.Resolver.Resolve(ctx, uid, meter)
	if err != nil {
		return "", err
	}
	if !has {
		return "", domain.ErrNoKey
	}
	if mode == domain.KeyModePlatform {
		if s.Gate != nil && !s.Gate.AllowPlatform(meter) {
			return "", domain.ErrPlatformCapacity
		}
		if err := s.reservePlatform(ctx, uid, meter); err != nil {
			return "", err
		}
	}
	return mode, nil
}

// MarkPlatformLimited cools down platform keys (e.g. after upstream 429).
// BYOK users are unaffected. until zero clears the cool-down.
func (s *Service) MarkPlatformLimited(until time.Time) {
	if s == nil || s.Gate == nil {
		return
	}
	s.Gate.MarkLimited(until)
}

func (s *Service) reservePlatform(ctx context.Context, uid int64, meter string) error {
	prefs, err := s.ensurePrefs(ctx, uid)
	if err != nil {
		return err
	}
	plan := domain.ResolvePlan(prefs.PlanID)
	totalLimit, meterLimit := plan.MonthlyCredits, -1
	if capN, ok := plan.SoftCaps[meter]; ok && capN > 0 {
		meterLimit = capN
	}
	if prefs.Unlimited {
		totalLimit, meterLimit = -1, -1
	}
	return s.Repo.ReservePlatform(ctx, uid, domain.CurrentMonthKey(), meter, domain.MeterCost(meter), totalLimit, meterLimit)
}

func (s *Service) ensurePrefs(ctx context.Context, uid int64) (*domain.MemberPrefs, error) {
	p, err := s.Repo.GetPrefs(ctx, uid)
	if err != nil {
		return nil, err
	}
	if p == nil {
		p = &domain.MemberPrefs{UID: uid, PlanID: domain.PlanFree, UpdatedAt: domain.NowNano()}
		if err := s.Repo.SavePrefs(ctx, p); err != nil {
			return nil, err
		}
	}
	// 正規化 plan_id（小寫／有效方案）
	norm := domain.ResolvePlan(p.PlanID).ID
	if p.PlanID != norm {
		p.PlanID = norm
		p.UpdatedAt = domain.NowNano()
		if err := s.Repo.SavePrefs(ctx, p); err != nil {
			return nil, err
		}
	}
	return p, nil
}

func (s *Service) GetSummary(ctx context.Context, uid int64, monthKey string) (*domain.MonthSummary, error) {
	if monthKey == "" {
		monthKey = domain.CurrentMonthKey()
	}
	prefs, err := s.ensurePrefs(ctx, uid)
	if err != nil {
		return nil, err
	}
	plan := domain.ResolvePlan(prefs.PlanID)
	// 回寫正規化 plan_id，避免 DB 殘留大小寫／空白導致永遠 fallback Free
	if prefs.PlanID != plan.ID {
		prefs.PlanID = plan.ID
		prefs.UpdatedAt = domain.NowNano()
		_ = s.Repo.SavePrefs(ctx, prefs)
	}
	events, err := s.Repo.ListEvents(ctx, uid, monthKey, "all", 0)
	if err != nil {
		return nil, err
	}
	platBy := map[string]domain.MeterCount{}
	byokBy := map[string]domain.MeterCount{}
	platUsed, platCalls, byokCalls := 0, 0, 0
	for _, e := range events {
		if e.KeyMode == domain.KeyModePlatform {
			platUsed += e.Credits
			platCalls++
			m := platBy[e.Meter]
			m.Count++
			m.Credits += e.Credits
			platBy[e.Meter] = m
		} else if e.KeyMode == domain.KeyModeByok {
			byokCalls++
			m := byokBy[e.Meter]
			m.Count++
			byokBy[e.Meter] = m
		}
	}
	if counter, counterErr := s.Repo.GetMonthlyCounter(ctx, uid, monthKey); counterErr == nil {
		platUsed = counter.TotalCredits
		for meter, credits := range counter.ByMeter {
			m := platBy[meter]
			m.Credits = credits
			platBy[meter] = m
		}
	} else if counterErr != domain.ErrNotFound {
		return nil, counterErr
	}
	remaining := plan.MonthlyCredits - platUsed
	if remaining < 0 {
		remaining = 0
	}
	// Unlimited 只影響是否擋用，顯示仍報真實已用／剩餘（前端以 unlimited 旗標顯示 ∞）
	return &domain.MonthSummary{
		MonthKey: monthKey, UID: uid, PlanID: plan.ID, Unlimited: prefs.Unlimited,
		Platform: domain.PlatformBlock{
			CreditsUsed: platUsed, CreditsRemaining: remaining, CreditsTotal: plan.MonthlyCredits,
			CallCount: platCalls, ByMeter: platBy,
		},
		Byok: domain.ByokBlock{CallCount: byokCalls, ByMeter: byokBy},
		// TotalCredits = 方案配給（legacy 欄位名）；已用請看 platform.credits_used
		TotalCredits: plan.MonthlyCredits, RemainingCredits: remaining,
	}, nil
}

func (s *Service) ListEvents(ctx context.Context, uid int64, monthKey, keyMode string, limit int) ([]*domain.Event, error) {
	if limit <= 0 {
		limit = 100
	}
	return s.Repo.ListEvents(ctx, uid, monthKey, keyMode, limit)
}

func (s *Service) GetPrefs(ctx context.Context, uid int64) (*domain.MemberPrefs, error) {
	return s.ensurePrefs(ctx, uid)
}

func (s *Service) SetPrefs(ctx context.Context, actorUID int64, targetUID int64, isAdmin bool, planID string, unlimited *bool) (*domain.MemberPrefs, error) {
	if !isAdmin && actorUID != targetUID {
		return nil, domain.ErrForbidden
	}
	if !isAdmin && unlimited != nil {
		// members cannot set unlimited
		return nil, domain.ErrForbidden
	}
	p, err := s.ensurePrefs(ctx, targetUID)
	if err != nil {
		return nil, err
	}
	if planID != "" {
		id := strings.ToLower(strings.TrimSpace(planID))
		if _, ok := domain.Plans[id]; !ok {
			return nil, domain.ErrInvalidPlan
		}
		// member self can only change via purchase
		if !isAdmin && actorUID == targetUID {
			return nil, domain.ErrForbidden
		}
		p.PlanID = id
		p.AdminOverride = true
	}
	if unlimited != nil && isAdmin {
		p.Unlimited = *unlimited
	}
	p.UpdatedAt = domain.NowNano()
	if err := s.Repo.SavePrefs(ctx, p); err != nil {
		return nil, err
	}
	return p, nil
}

func (s *Service) PurchasePlan(ctx context.Context, uid int64, planID, mockRef string) (*domain.Purchase, error) {
	return nil, domain.ErrPurchaseDisabled
}

func (s *Service) ListMyPurchases(ctx context.Context, uid int64, limit int) ([]*domain.Purchase, error) {
	if limit <= 0 {
		limit = 50
	}
	return s.Repo.ListPurchases(ctx, uid, limit)
}

func (s *Service) TenantSummary(ctx context.Context, monthKey string) (*domain.TenantSummary, error) {
	if monthKey == "" {
		monthKey = domain.CurrentMonthKey()
	}
	events, err := s.Repo.ListEventsAll(ctx, monthKey, 0)
	if err != nil {
		return nil, err
	}
	type acc struct {
		plat int
		byok int
	}
	byUID := map[int64]*acc{}
	platTotal, byokTotal := 0, 0
	for _, e := range events {
		a := byUID[e.UID]
		if a == nil {
			a = &acc{}
			byUID[e.UID] = a
		}
		if e.KeyMode == domain.KeyModePlatform {
			a.plat += e.Credits
			platTotal += e.Credits
		} else if e.KeyMode == domain.KeyModeByok {
			a.byok++
			byokTotal++
		}
	}
	var members []domain.TenantRow
	for uid, a := range byUID {
		prefs, err := s.ensurePrefs(ctx, uid)
		if err != nil {
			return nil, err
		}
		members = append(members, domain.TenantRow{
			UID: uid, PlanID: prefs.PlanID, Unlimited: prefs.Unlimited,
			PlatformCreditsUsed: a.plat, ByokCallCount: a.byok,
		})
	}
	return &domain.TenantSummary{
		MonthKey: monthKey, PlatformCreditsUsed: platTotal, ByokCallCount: byokTotal, Members: members,
	}, nil
}