geo/server/internal/tenant/repository/quota_repo.go

package repository

import (
	"context"
	"errors"
	"time"

	"github.com/jackc/pgx/v5"

	"github.com/geo-platform/tenant-api/internal/tenant/repository/generated"
)

type ArticleQuotaStatus struct {
	PlanCode string
	PlanName string
	Total    int
	Used     int
	Balance  int
	ResetAt  *time.Time
}

type AIQuotaStatus struct {
	PlanCode  string
	PlanName  string
	Total     int
	Used      int
	Balance   int
	BaseChars int
	ResetAt   *time.Time
}

type QuotaLedgerInput struct {
	TenantID      int64
	OperatorID    int64
	QuotaType     string
	Delta         int
	BalanceAfter  int
	Reason        *string
	ReferenceType *string
	ReferenceID   *int64
}

type QuotaReservationInput struct {
	TenantID       int64
	OperatorID     int64
	QuotaType      string
	ResourceType   string
	ResourceID     *int64
	ReservedAmount int
	ExpireAt       time.Time
}

type QuotaRepository interface {
	GetCurrentBalance(ctx context.Context, tenantID int64, quotaType string) (int, error)
	GetArticleQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*ArticleQuotaStatus, error)
	GetAIQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*AIQuotaStatus, error)
	InsertQuotaLedger(ctx context.Context, input QuotaLedgerInput) (int64, error)
	CreateQuotaReservation(ctx context.Context, input QuotaReservationInput) (int64, error)
	UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error
	ConfirmReservation(ctx context.Context, reservationID, tenantID int64) error
	RefundReservation(ctx context.Context, reservationID, tenantID int64) error
}

type quotaRepository struct {
	q  generated.Querier
	db generated.DBTX
}

func NewQuotaRepository(db generated.DBTX) QuotaRepository {
	return &quotaRepository{q: newQuerier(db), db: db}
}

func (r *quotaRepository) GetCurrentBalance(ctx context.Context, tenantID int64, quotaType string) (int, error) {
	if quotaType != "article_generation" {
		balance, err := r.q.GetCurrentBalance(ctx, generated.GetCurrentBalanceParams{
			TenantID:  tenantID,
			QuotaType: quotaType,
		})
		if err != nil {
			return 0, err
		}
		return int(balance), nil
	}

	status, err := r.GetArticleQuotaStatus(ctx, tenantID, time.Now().UTC())
	if err != nil {
		return 0, err
	}
	return status.Balance, nil
}

func (r *quotaRepository) GetAIQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*AIQuotaStatus, error) {
	access, err := loadTenantPlanAccess(ctx, r.db, tenantID, now)
	if err != nil {
		return nil, err
	}
	if access == nil {
		return nil, pgx.ErrNoRows
	}

	baseChars := access.AIPointBaseChars()
	if !access.HasActiveAccess(now) {
		return &AIQuotaStatus{
			PlanCode:  access.PlanCode,
			PlanName:  access.PlanName,
			BaseChars: baseChars,
		}, nil
	}

	total := access.AIPointsMonthlyLimit()
	if total <= 0 {
		return &AIQuotaStatus{
			PlanCode:  access.PlanCode,
			PlanName:  access.PlanName,
			BaseChars: baseChars,
		}, nil
	}

	windowStart, windowEnd, resetAt := access.AIPointsWindow(now)
	if windowEnd.IsZero() {
		return nil, errors.New("ai points quota window end is missing")
	}

	used, err := CountEffectiveQuotaReservations(ctx, r.db, tenantID, "ai_points", windowStart, windowEnd)
	if err != nil {
		return nil, err
	}
	balance := total - used
	if balance < 0 {
		balance = 0
	}

	return &AIQuotaStatus{
		PlanCode:  access.PlanCode,
		PlanName:  access.PlanName,
		Total:     total,
		Used:      used,
		Balance:   balance,
		BaseChars: baseChars,
		ResetAt:   resetAt,
	}, nil
}

func (r *quotaRepository) GetArticleQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*ArticleQuotaStatus, error) {
	access, err := loadTenantPlanAccess(ctx, r.db, tenantID, now)
	if err != nil {
		return nil, err
	}
	if access == nil {
		return nil, pgx.ErrNoRows
	}

	// Non-HTTP callers (scheduler, workers) bypass subscription_guard, so any
	// code path that consults article quota must itself enforce that the
	// subscription is currently active. Return a zero-balance status when the
	// plan is expired/inactive/not-yet-active so the caller short-circuits on
	// insufficient balance instead of dispatching generation.
	if !access.HasActiveAccess(now) {
		return &ArticleQuotaStatus{
			PlanCode: access.PlanCode,
			PlanName: access.PlanName,
			Total:    0,
			Used:     0,
			Balance:  0,
		}, nil
	}

	total := access.ArticleGenerationLimit()
	if total <= 0 {
		return &ArticleQuotaStatus{
			PlanCode: access.PlanCode,
			PlanName: access.PlanName,
			Total:    0,
			Used:     0,
			Balance:  0,
		}, nil
	}

	windowStart, windowEnd, resetAt := access.ArticleQuotaWindow(now)
	if windowEnd.IsZero() {
		return nil, errors.New("article quota window end is missing")
	}

	used, err := CountEffectiveQuotaReservations(ctx, r.db, tenantID, "article_generation", windowStart, windowEnd)
	if err != nil {
		return nil, err
	}
	balance := total - used
	if balance < 0 {
		balance = 0
	}

	return &ArticleQuotaStatus{
		PlanCode: access.PlanCode,
		PlanName: access.PlanName,
		Total:    total,
		Used:     used,
		Balance:  balance,
		ResetAt:  resetAt,
	}, nil
}

func (r *quotaRepository) InsertQuotaLedger(ctx context.Context, input QuotaLedgerInput) (int64, error) {
	var operatorID *int64
	if input.OperatorID > 0 {
		operatorID = &input.OperatorID
	}
	return r.q.InsertQuotaLedger(ctx, generated.InsertQuotaLedgerParams{
		TenantID:      input.TenantID,
		OperatorID:    pgInt8(operatorID),
		QuotaType:     input.QuotaType,
		Delta:         int32(input.Delta),
		BalanceAfter:  int32(input.BalanceAfter),
		Reason:        pgText(input.Reason),
		ReferenceType: pgText(input.ReferenceType),
		ReferenceID:   pgInt8(input.ReferenceID),
	})
}

func (r *quotaRepository) CreateQuotaReservation(ctx context.Context, input QuotaReservationInput) (int64, error) {
	expireAt := input.ExpireAt
	var operatorID *int64
	if input.OperatorID > 0 {
		operatorID = &input.OperatorID
	}
	return r.q.CreateQuotaReservation(ctx, generated.CreateQuotaReservationParams{
		TenantID:       input.TenantID,
		OperatorID:     pgInt8(operatorID),
		QuotaType:      input.QuotaType,
		ResourceType:   input.ResourceType,
		ResourceID:     pgInt8(input.ResourceID),
		ReservedAmount: int32(input.ReservedAmount),
		ExpireAt:       pgTimestamp(&expireAt),
	})
}

func (r *quotaRepository) UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error {
	resourceIDValue := resourceID
	return r.q.UpdateQuotaReservationResource(ctx, generated.UpdateQuotaReservationResourceParams{
		ResourceID: pgInt8(&resourceIDValue),
		ID:         reservationID,
		TenantID:   tenantID,
	})
}

func (r *quotaRepository) ConfirmReservation(ctx context.Context, reservationID, tenantID int64) error {
	return r.q.ConfirmReservation(ctx, generated.ConfirmReservationParams{
		ID:       reservationID,
		TenantID: tenantID,
	})
}

func (r *quotaRepository) RefundReservation(ctx context.Context, reservationID, tenantID int64) error {
	return r.q.RefundReservation(ctx, generated.RefundReservationParams{
		ID:       reservationID,
		TenantID: tenantID,
	})
}