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

package repository

import (
	"context"
	"errors"
	"time"

	"github.com/jackc/pgx/v5"

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

type KolSubscription struct {
	ID         int64
	TenantID   int64
	PackageID  int64
	Status     string
	ApprovedBy *int64
	StartAt    *time.Time
	EndAt      *time.Time
	CreatedAt  time.Time
	UpdatedAt  time.Time
}

type KolSubscriptionPrompt struct {
	ID                  int64
	TenantID            int64
	SubscriptionID      int64
	PackageID           int64
	PromptID            int64
	Status              string
	GrantedAt           *time.Time
	RevokedAt           *time.Time
	PromptName          string
	PlatformHint        *string
	PackageName         string
	PublishedRevisionNo *int
	// KolDisplayName is populated by ListSubscriptionPromptsByTenant (joins kol_profiles).
	KolDisplayName string
	// Fields below only populated by GetSubscriptionPromptByID.
	SubscriptionStatus string
	SubscriptionEndAt  *time.Time
	PromptStatus       string
	PackageStatus      string
}

type CreateAccessRowInput struct {
	TenantID       int64
	SubscriptionID int64
	PackageID      int64
	PromptID       int64
}

type KolSubscriptionRepository interface {
	Create(ctx context.Context, tenantID, packageID int64) (*KolSubscription, error)
	GetByID(ctx context.Context, tenantID, id int64) (*KolSubscription, error)
	GetActiveForTenant(ctx context.Context, tenantID, packageID int64) (*KolSubscription, error)
	ListByTenant(ctx context.Context, tenantID int64) ([]KolSubscription, error)
	ListActiveSubscribersForPackage(ctx context.Context, packageID int64) ([]KolSubscription, error)
	Approve(ctx context.Context, id, tenantID, approvedBy int64, endAt *time.Time) error
	RevokeByTenant(ctx context.Context, id, tenantID int64) error
	CreateAccessRow(ctx context.Context, input CreateAccessRowInput) error
	ListSubscriptionPromptsByTenant(ctx context.Context, tenantID int64) ([]KolSubscriptionPrompt, error)
	GetSubscriptionPromptByID(ctx context.Context, tenantID, id int64) (*KolSubscriptionPrompt, error)
	RevokeAccessRowsBySubscription(ctx context.Context, subscriptionID, tenantID int64) error
	RevokeAccessRowsByPrompt(ctx context.Context, promptID, tenantID int64) error
}

type kolSubscriptionRepository struct {
	q generated.Querier
}

func NewKolSubscriptionRepository(db generated.DBTX) KolSubscriptionRepository {
	return &kolSubscriptionRepository{q: newQuerier(db)}
}

func mapKolSubscription(row generated.KolSubscription) KolSubscription {
	startAt := optionalTime(row.StartAt)
	endAt := optionalTime(row.EndAt)
	return KolSubscription{
		ID:         row.ID,
		TenantID:   row.TenantID,
		PackageID:  row.PackageID,
		Status:     row.Status,
		ApprovedBy: nullableInt64(row.ApprovedBy),
		StartAt:    startAt,
		EndAt:      endAt,
		CreatedAt:  timeFromTimestamp(row.CreatedAt),
		UpdatedAt:  timeFromTimestamp(row.UpdatedAt),
	}
}

func (r *kolSubscriptionRepository) Create(ctx context.Context, tenantID, packageID int64) (*KolSubscription, error) {
	row, err := r.q.CreateKolSubscription(ctx, generated.CreateKolSubscriptionParams{
		TenantID:  tenantID,
		PackageID: packageID,
	})
	if err != nil {
		return nil, err
	}
	sub := mapKolSubscription(row)
	return &sub, nil
}

func (r *kolSubscriptionRepository) GetByID(ctx context.Context, tenantID, id int64) (*KolSubscription, error) {
	row, err := r.q.GetKolSubscriptionByID(ctx, generated.GetKolSubscriptionByIDParams{
		ID:       id,
		TenantID: tenantID,
	})
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	sub := mapKolSubscription(row)
	return &sub, nil
}

func (r *kolSubscriptionRepository) GetActiveForTenant(ctx context.Context, tenantID, packageID int64) (*KolSubscription, error) {
	row, err := r.q.GetActiveKolSubscription(ctx, generated.GetActiveKolSubscriptionParams{
		TenantID:  tenantID,
		PackageID: packageID,
	})
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	sub := mapKolSubscription(row)
	return &sub, nil
}

func (r *kolSubscriptionRepository) ListByTenant(ctx context.Context, tenantID int64) ([]KolSubscription, error) {
	rows, err := r.q.ListKolSubscriptions(ctx, tenantID)
	if err != nil {
		return nil, err
	}
	result := make([]KolSubscription, len(rows))
	for i, row := range rows {
		result[i] = mapKolSubscription(row)
	}
	return result, nil
}

func (r *kolSubscriptionRepository) ListActiveSubscribersForPackage(ctx context.Context, packageID int64) ([]KolSubscription, error) {
	rows, err := r.q.ListActiveSubscribersForPackage(ctx, packageID)
	if err != nil {
		return nil, err
	}
	result := make([]KolSubscription, len(rows))
	for i, row := range rows {
		result[i] = mapKolSubscription(row)
	}
	return result, nil
}

func (r *kolSubscriptionRepository) Approve(ctx context.Context, id, tenantID, approvedBy int64, endAt *time.Time) error {
	return r.q.ApproveKolSubscription(ctx, generated.ApproveKolSubscriptionParams{
		ApprovedBy: pgInt8(&approvedBy),
		EndAt:      pgTimestamp(endAt),
		ID:         id,
		TenantID:   tenantID,
	})
}

func (r *kolSubscriptionRepository) RevokeByTenant(ctx context.Context, id, tenantID int64) error {
	return r.q.RevokeKolSubscription(ctx, generated.RevokeKolSubscriptionParams{
		ID:       id,
		TenantID: tenantID,
	})
}

func (r *kolSubscriptionRepository) CreateAccessRow(ctx context.Context, input CreateAccessRowInput) error {
	return r.q.CreateSubscriptionPrompt(ctx, generated.CreateSubscriptionPromptParams{
		TenantID:       input.TenantID,
		SubscriptionID: input.SubscriptionID,
		PackageID:      input.PackageID,
		PromptID:       input.PromptID,
	})
}

func (r *kolSubscriptionRepository) ListSubscriptionPromptsByTenant(ctx context.Context, tenantID int64) ([]KolSubscriptionPrompt, error) {
	rows, err := r.q.ListSubscriptionPromptsByTenant(ctx, tenantID)
	if err != nil {
		return nil, err
	}
	result := make([]KolSubscriptionPrompt, len(rows))
	for i, row := range rows {
		grantedAt := optionalTime(row.GrantedAt)
		revokedAt := optionalTime(row.RevokedAt)
		result[i] = KolSubscriptionPrompt{
			ID:                  row.ID,
			TenantID:            row.TenantID,
			SubscriptionID:      row.SubscriptionID,
			PackageID:           row.PackageID,
			PromptID:            row.PromptID,
			Status:              row.Status,
			GrantedAt:           grantedAt,
			RevokedAt:           revokedAt,
			PromptName:          row.PromptName,
			PlatformHint:        nullableText(row.PlatformHint),
			PackageName:         row.PackageName,
			PublishedRevisionNo: nullableIntFromInt4(row.PublishedRevisionNo),
			KolDisplayName:      row.KolDisplayName,
		}
	}
	return result, nil
}

func (r *kolSubscriptionRepository) GetSubscriptionPromptByID(ctx context.Context, tenantID, id int64) (*KolSubscriptionPrompt, error) {
	row, err := r.q.GetSubscriptionPromptByID(ctx, generated.GetSubscriptionPromptByIDParams{
		ID:       id,
		TenantID: tenantID,
	})
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	grantedAt := optionalTime(row.GrantedAt)
	revokedAt := optionalTime(row.RevokedAt)
	subEndAt := optionalTime(row.SubscriptionEndAt)
	return &KolSubscriptionPrompt{
		ID:                  row.ID,
		TenantID:            row.TenantID,
		SubscriptionID:      row.SubscriptionID,
		PackageID:           row.PackageID,
		PromptID:            row.PromptID,
		Status:              row.Status,
		GrantedAt:           grantedAt,
		RevokedAt:           revokedAt,
		PromptName:          row.PromptName,
		PlatformHint:        nullableText(row.PlatformHint),
		PackageName:         row.PackageName,
		PublishedRevisionNo: nullableIntFromInt4(row.PublishedRevisionNo),
		SubscriptionStatus:  row.SubscriptionStatus,
		SubscriptionEndAt:   subEndAt,
		PromptStatus:        row.PromptStatus,
		PackageStatus:       row.PackageStatus,
	}, nil
}

func (r *kolSubscriptionRepository) RevokeAccessRowsBySubscription(ctx context.Context, subscriptionID, tenantID int64) error {
	return r.q.RevokeSubscriptionPromptsBySubscription(ctx, generated.RevokeSubscriptionPromptsBySubscriptionParams{
		SubscriptionID: subscriptionID,
		TenantID:       tenantID,
	})
}

func (r *kolSubscriptionRepository) RevokeAccessRowsByPrompt(ctx context.Context, promptID, tenantID int64) error {
	return r.q.RevokeSubscriptionPromptsByPrompt(ctx, generated.RevokeSubscriptionPromptsByPromptParams{
		PromptID: promptID,
		TenantID: tenantID,
	})
}