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

package repository

import (
	"context"
	"errors"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"

	"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
	CreatorTenantID    int64
}

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

type KolSubscriptionRepository interface {
	Create(ctx context.Context, tenantID, packageID int64) (*KolSubscription, error)
	SelfSubscribe(ctx context.Context, tenantID, packageID int64) (*KolSubscription, error)
	CancelSelfSubscription(ctx context.Context, tenantID, packageID int64) (*KolSubscription, error)
	GetByID(ctx context.Context, tenantID, id int64) (*KolSubscription, error)
	GetByIDAny(ctx context.Context, 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 {
	db generated.DBTX
	q  generated.Querier
}

func NewKolSubscriptionRepository(db generated.DBTX) KolSubscriptionRepository {
	return &kolSubscriptionRepository{db: db, 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) SelfSubscribe(ctx context.Context, tenantID, packageID int64) (*KolSubscription, error) {
	tx, err := beginKolSubscriptionTx(ctx, r.db)
	if err != nil {
		return nil, err
	}
	if tx != nil {
		defer tx.Rollback(ctx)
	}
	db := kolSubscriptionDB(r.db, tx)

	sub, err := lockCurrentKolSubscription(ctx, db, tenantID, packageID)
	if err != nil {
		return nil, err
	}

	if sub == nil {
		sub, err = insertSelfKolSubscription(ctx, db, tenantID, packageID)
	} else if sub.Status != "active" {
		sub, err = activateSelfKolSubscription(ctx, db, sub.ID, tenantID)
	}
	if err != nil {
		return nil, err
	}

	if err := grantSelfSubscriptionPrompts(ctx, db, tenantID, packageID, sub.ID); err != nil {
		return nil, err
	}

	updated, err := getKolSubscriptionByTenantAndID(ctx, db, tenantID, sub.ID)
	if err != nil {
		return nil, err
	}

	if tx != nil {
		if err := tx.Commit(ctx); err != nil {
			return nil, err
		}
	}
	return updated, nil
}

func (r *kolSubscriptionRepository) CancelSelfSubscription(ctx context.Context, tenantID, packageID int64) (*KolSubscription, error) {
	tx, err := beginKolSubscriptionTx(ctx, r.db)
	if err != nil {
		return nil, err
	}
	if tx != nil {
		defer tx.Rollback(ctx)
	}
	db := kolSubscriptionDB(r.db, tx)

	sub, err := lockCurrentKolSubscription(ctx, db, tenantID, packageID)
	if err != nil {
		return nil, err
	}
	if sub == nil {
		return nil, nil
	}

	updated, err := revokeKolSubscriptionByID(ctx, db, tenantID, sub.ID)
	if err != nil {
		return nil, err
	}
	if err := revokeKolSubscriptionAccessRows(ctx, db, tenantID, sub.ID); err != nil {
		return nil, err
	}

	if tx != nil {
		if err := tx.Commit(ctx); err != nil {
			return nil, err
		}
	}
	return updated, nil
}

func beginKolSubscriptionTx(ctx context.Context, db generated.DBTX) (pgx.Tx, error) {
	if starter, ok := db.(interface {
		Begin(context.Context) (pgx.Tx, error)
	}); ok {
		return starter.Begin(ctx)
	}
	return nil, nil
}

func kolSubscriptionDB(fallback generated.DBTX, tx pgx.Tx) generated.DBTX {
	if tx != nil {
		return tx
	}
	return fallback
}

func scanKolSubscriptionRaw(row interface{ Scan(...any) error }) (*KolSubscription, error) {
	var (
		id         int64
		tenantID   int64
		packageID  int64
		status     string
		approvedBy pgtype.Int8
		startAt    pgtype.Timestamptz
		endAt      pgtype.Timestamptz
		createdAt  pgtype.Timestamptz
		updatedAt  pgtype.Timestamptz
	)
	if err := row.Scan(
		&id,
		&tenantID,
		&packageID,
		&status,
		&approvedBy,
		&startAt,
		&endAt,
		&createdAt,
		&updatedAt,
	); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	return &KolSubscription{
		ID:         id,
		TenantID:   tenantID,
		PackageID:  packageID,
		Status:     status,
		ApprovedBy: nullableInt64(approvedBy),
		StartAt:    optionalTime(startAt),
		EndAt:      optionalTime(endAt),
		CreatedAt:  timeFromTimestamp(createdAt),
		UpdatedAt:  timeFromTimestamp(updatedAt),
	}, nil
}

func lockCurrentKolSubscription(ctx context.Context, db generated.DBTX, tenantID, packageID int64) (*KolSubscription, error) {
	return scanKolSubscriptionRaw(db.QueryRow(ctx, `
SELECT id, tenant_id, package_id, status, approved_by, start_at, end_at, created_at, updated_at
FROM kol_subscriptions
WHERE tenant_id = $1
  AND package_id = $2
  AND status IN ('pending', 'active')
ORDER BY created_at DESC, id DESC
LIMIT 1
FOR UPDATE`, tenantID, packageID))
}

func insertSelfKolSubscription(ctx context.Context, db generated.DBTX, tenantID, packageID int64) (*KolSubscription, error) {
	return scanKolSubscriptionRaw(db.QueryRow(ctx, `
INSERT INTO kol_subscriptions (tenant_id, package_id, status, approved_by, start_at, end_at)
VALUES ($1, $2, 'active', NULL, NOW(), NULL)
RETURNING id, tenant_id, package_id, status, approved_by, start_at, end_at, created_at, updated_at`, tenantID, packageID))
}

func activateSelfKolSubscription(ctx context.Context, db generated.DBTX, id, tenantID int64) (*KolSubscription, error) {
	return scanKolSubscriptionRaw(db.QueryRow(ctx, `
UPDATE kol_subscriptions
SET status = 'active',
    approved_by = NULL,
    start_at = NOW(),
    end_at = NULL,
    updated_at = NOW()
WHERE id = $1
  AND tenant_id = $2
RETURNING id, tenant_id, package_id, status, approved_by, start_at, end_at, created_at, updated_at`, id, tenantID))
}

func getKolSubscriptionByTenantAndID(ctx context.Context, db generated.DBTX, tenantID, id int64) (*KolSubscription, error) {
	return scanKolSubscriptionRaw(db.QueryRow(ctx, `
SELECT id, tenant_id, package_id, status, approved_by, start_at, end_at, created_at, updated_at
FROM kol_subscriptions
WHERE id = $1
  AND tenant_id = $2`, id, tenantID))
}

func revokeKolSubscriptionByID(ctx context.Context, db generated.DBTX, tenantID, id int64) (*KolSubscription, error) {
	return scanKolSubscriptionRaw(db.QueryRow(ctx, `
UPDATE kol_subscriptions
SET status = 'revoked',
    updated_at = NOW()
WHERE id = $1
  AND tenant_id = $2
  AND status IN ('pending', 'active')
RETURNING id, tenant_id, package_id, status, approved_by, start_at, end_at, created_at, updated_at`, id, tenantID))
}

func grantSelfSubscriptionPrompts(ctx context.Context, db generated.DBTX, tenantID, packageID, subscriptionID int64) error {
	_, err := db.Exec(ctx, `
INSERT INTO kol_subscription_prompts (tenant_id, subscription_id, package_id, prompt_id, status)
SELECT $1, $2, $3, p.id, 'active'
FROM kol_prompts p
WHERE p.tenant_id = $1
  AND p.package_id = $3
  AND p.status = 'active'
  AND p.published_revision_no IS NOT NULL
  AND p.deleted_at IS NULL
ON CONFLICT (subscription_id, prompt_id) DO UPDATE
SET status = 'active',
    revoked_at = NULL,
    updated_at = NOW()`, tenantID, subscriptionID, packageID)
	return err
}

func revokeKolSubscriptionAccessRows(ctx context.Context, db generated.DBTX, tenantID, subscriptionID int64) error {
	_, err := db.Exec(ctx, `
UPDATE kol_subscription_prompts
SET status = 'revoked',
    revoked_at = NOW(),
    updated_at = NOW()
WHERE tenant_id = $1
  AND subscription_id = $2
  AND status <> 'revoked'`, tenantID, subscriptionID)
	return err
}

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) GetByIDAny(ctx context.Context, id int64) (*KolSubscription, error) {
	row, err := r.q.GetKolSubscriptionByIDAny(ctx, id)
	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,
		CreatorTenantID:     row.CreatorTenantID,
	}, 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,
	})
}