geo/server/internal/tenant/app/ai_point_usage.go

package app

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"time"
	"unicode/utf8"

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

	sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
	"github.com/geo-platform/tenant-api/internal/shared/response"
	"github.com/geo-platform/tenant-api/internal/tenant/repository"
)

const (
	AIUsageTypeArticleSelectionOptimize = "article_selection_optimize"
	AIUsageTypeTemplateAnalyze          = "template_analyze"
	AIUsageTypeTemplateTitleGenerate    = "template_title_generate"
	AIUsageTypeTemplateOutlineGenerate  = "template_outline_generate"
	AIUsageTypeKolPromptGenerate        = "kol_prompt_generate"
	AIUsageTypeKolPromptOptimize        = "kol_prompt_optimize"
	AIUsageTypeQuestionDistill          = "question_distill"
	AIUsageTypeQuestionCombinationFill  = "question_combination_fill"

	aiPointsQuotaType = "ai_points"
)

type AIPointReserveInput struct {
	TenantID     int64
	OperatorID   int64
	UsageType    string
	ResourceType *string
	ResourceID   *int64
	ResourceUID  *string
	MeteredText  string
	FixedPoints  int
	Metadata     map[string]any
}

type AIPointReservation struct {
	ReservationID int64
	UsageLogID    int64
	Points        int
	RequestChars  int
	BaseChars     int
}

func ReserveAIPoints(
	ctx context.Context,
	pool *pgxpool.Pool,
	cache sharedcache.Cache,
	input AIPointReserveInput,
) (*AIPointReservation, error) {
	if pool == nil {
		return nil, response.ErrServiceUnavailable(50344, "ai_points_unavailable", "ai points store is unavailable")
	}

	input.UsageType = strings.TrimSpace(input.UsageType)
	if input.UsageType == "" {
		return nil, response.ErrBadRequest(40075, "ai_points_usage_type_required", "usage type is required")
	}

	tx, err := pool.Begin(ctx)
	if err != nil {
		return nil, response.ErrInternal(50120, "ai_points_tx_failed", "failed to begin ai points transaction")
	}
	defer rollbackGenerationTx(tx)

	if err := lockTenantAIPoints(ctx, tx, input.TenantID); err != nil {
		return nil, response.ErrInternal(50121, "ai_points_lock_failed", "failed to lock ai points quota")
	}

	now := time.Now().UTC()
	quotaRepo := repository.NewQuotaRepository(tx)
	status, err := quotaRepo.GetAIQuotaStatus(ctx, input.TenantID, now)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are not available for this plan")
		}
		return nil, response.ErrInternal(50122, "ai_points_query_failed", "failed to query ai points")
	}

	requestChars := countAIPointChars(input.MeteredText)
	points := calculateAIPointCost(requestChars, status.BaseChars)
	if input.FixedPoints > 0 {
		points = input.FixedPoints
	}
	if status.Total <= 0 || status.Balance < points {
		return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are insufficient, please upgrade your plan")
	}

	resourceType := "ai_usage"
	if input.ResourceType != nil && strings.TrimSpace(*input.ResourceType) != "" {
		resourceType = strings.TrimSpace(*input.ResourceType)
	}

	reservationID, err := quotaRepo.CreateQuotaReservation(ctx, repository.QuotaReservationInput{
		TenantID:       input.TenantID,
		OperatorID:     input.OperatorID,
		QuotaType:      aiPointsQuotaType,
		ResourceType:   resourceType,
		ResourceID:     input.ResourceID,
		ReservedAmount: points,
		ExpireAt:       now.Add(time.Hour),
	})
	if err != nil {
		return nil, response.ErrInternal(50123, "ai_points_reserve_failed", "failed to reserve ai points")
	}

	metadata := input.Metadata
	if metadata == nil {
		metadata = map[string]any{}
	}
	metadataJSON, err := json.Marshal(metadata)
	if err != nil {
		return nil, response.ErrInternal(50124, "ai_points_metadata_failed", "failed to encode ai points metadata")
	}
	usageID, err := repository.NewAIPointUsageRepository(tx).Create(ctx, repository.CreateAIPointUsageLogInput{
		TenantID:           input.TenantID,
		OperatorID:         input.OperatorID,
		QuotaReservationID: reservationID,
		UsageType:          input.UsageType,
		ResourceType:       &resourceType,
		ResourceID:         input.ResourceID,
		ResourceUID:        input.ResourceUID,
		RequestChars:       requestChars,
		BaseChars:          status.BaseChars,
		Points:             points,
		MetadataJSON:       metadataJSON,
	})
	if err != nil {
		return nil, response.ErrInternal(50125, "ai_points_audit_failed", "failed to create ai points audit log")
	}

	reason := "ai_points_reserve"
	referenceType := "ai_point_usage"
	balanceAfter := status.Balance - points
	if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
		TenantID:      input.TenantID,
		OperatorID:    input.OperatorID,
		QuotaType:     aiPointsQuotaType,
		Delta:         -points,
		BalanceAfter:  balanceAfter,
		Reason:        &reason,
		ReferenceType: &referenceType,
		ReferenceID:   &usageID,
	}); err != nil {
		return nil, response.ErrInternal(50126, "ai_points_ledger_failed", "failed to write ai points ledger")
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, response.ErrInternal(50127, "ai_points_commit_failed", "failed to commit ai points reservation")
	}
	deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(input.TenantID))

	return &AIPointReservation{
		ReservationID: reservationID,
		UsageLogID:    usageID,
		Points:        points,
		RequestChars:  requestChars,
		BaseChars:     status.BaseChars,
	}, nil
}

func CompleteAIPoints(
	ctx context.Context,
	pool *pgxpool.Pool,
	tenantID int64,
	reservation AIPointReservation,
	model string,
) error {
	if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 {
		return nil
	}

	tx, err := pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer rollbackGenerationTx(tx)

	pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID)
	if err != nil || !pending {
		return err
	}
	if err := confirmAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID); err != nil {
		return err
	}
	var modelPtr *string
	if strings.TrimSpace(model) != "" {
		normalized := strings.TrimSpace(model)
		modelPtr = &normalized
	}
	if err := repository.NewAIPointUsageRepository(tx).MarkCompleted(ctx, tenantID, reservation.UsageLogID, modelPtr); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func RefundAIPoints(
	ctx context.Context,
	pool *pgxpool.Pool,
	cache sharedcache.Cache,
	tenantID int64,
	operatorID int64,
	reservation AIPointReservation,
	errorMessage string,
) error {
	if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 || reservation.Points <= 0 {
		return nil
	}

	tx, err := pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer rollbackGenerationTx(tx)

	if err := lockTenantAIPoints(ctx, tx, tenantID); err != nil {
		return err
	}
	pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID)
	if err != nil || !pending {
		return err
	}

	refunded, err := refundAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID)
	if err != nil {
		return err
	}
	if refunded {
		status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
		if err != nil {
			return err
		}
		quotaRepo := repository.NewQuotaRepository(tx)
		reason := "ai_points_refund"
		referenceType := "ai_point_usage"
		usageID := reservation.UsageLogID
		if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
			TenantID:      tenantID,
			OperatorID:    operatorID,
			QuotaType:     aiPointsQuotaType,
			Delta:         reservation.Points,
			BalanceAfter:  status.Balance + reservation.Points,
			Reason:        &reason,
			ReferenceType: &referenceType,
			ReferenceID:   &usageID,
		}); err != nil {
			return err
		}
	}
	if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil {
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return err
	}
	deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(tenantID))
	return nil
}

func CompletePendingAIPointsByResource(
	ctx context.Context,
	pool *pgxpool.Pool,
	tenantID int64,
	resourceType string,
	resourceID *int64,
	resourceUID *string,
	model string,
) error {
	reservation, err := findPendingAIPointReservation(ctx, pool, tenantID, resourceType, resourceID, resourceUID)
	if err != nil || reservation == nil {
		return err
	}
	return CompleteAIPoints(ctx, pool, tenantID, *reservation, model)
}

func RefundPendingAIPointsByResource(
	ctx context.Context,
	pool *pgxpool.Pool,
	cache sharedcache.Cache,
	tenantID int64,
	resourceType string,
	resourceID *int64,
	resourceUID *string,
	errorMessage string,
) error {
	reservation, operatorID, err := findPendingAIPointReservationWithOperator(ctx, pool, tenantID, resourceType, resourceID, resourceUID)
	if err != nil || reservation == nil {
		return err
	}
	return RefundAIPoints(ctx, pool, cache, tenantID, operatorID, *reservation, errorMessage)
}

func findPendingAIPointReservation(
	ctx context.Context,
	pool *pgxpool.Pool,
	tenantID int64,
	resourceType string,
	resourceID *int64,
	resourceUID *string,
) (*AIPointReservation, error) {
	reservation, _, err := findPendingAIPointReservationWithOperator(ctx, pool, tenantID, resourceType, resourceID, resourceUID)
	return reservation, err
}

func findPendingAIPointReservationWithOperator(
	ctx context.Context,
	pool *pgxpool.Pool,
	tenantID int64,
	resourceType string,
	resourceID *int64,
	resourceUID *string,
) (*AIPointReservation, int64, error) {
	if pool == nil {
		return nil, 0, nil
	}
	log, err := repository.NewAIPointUsageRepository(pool).FindPendingByResource(ctx, tenantID, resourceType, resourceID, resourceUID)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, 0, nil
		}
		return nil, 0, err
	}
	if log == nil || log.QuotaReservationID == nil {
		return nil, 0, nil
	}
	operatorID := int64(0)
	if log.OperatorID != nil {
		operatorID = *log.OperatorID
	}
	return &AIPointReservation{
		ReservationID: *log.QuotaReservationID,
		UsageLogID:    log.ID,
		Points:        log.Points,
		RequestChars:  log.RequestChars,
		BaseChars:     log.BaseChars,
	}, operatorID, nil
}

func countAIPointChars(text string) int {
	return utf8.RuneCountInString(strings.TrimSpace(text))
}

func calculateAIPointCost(chars, baseChars int) int {
	if baseChars <= 0 {
		baseChars = 1000
	}
	if chars <= 0 {
		return 1
	}
	return chars/baseChars + 1
}

func lockTenantAIPoints(ctx context.Context, tx pgx.Tx, tenantID int64) error {
	_, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(20260428, hashint8($1)::int)`, tenantID)
	return err
}

func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogID int64) (bool, error) {
	var status string
	err := tx.QueryRow(ctx, `
	SELECT status
	FROM ai_point_usage_logs
	WHERE id = $1 AND tenant_id = $2
	FOR UPDATE
	`, usageLogID, tenantID).Scan(&status)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return false, nil
		}
		return false, err
	}
	return status == "pending", nil
}

func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
	tag, err := tx.Exec(ctx, `
	UPDATE quota_reservations
	SET status = 'confirmed',
	    consumed_amount = reserved_amount,
	    updated_at = NOW()
	WHERE id = $1
	  AND tenant_id = $2
	  AND quota_type = $3
	  AND status = 'pending'
	`, reservationID, tenantID, aiPointsQuotaType)
	if err != nil {
		return err
	}
	return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID)
}

func refundAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) (bool, error) {
	tag, err := tx.Exec(ctx, `
	UPDATE quota_reservations
	SET status = 'refunded',
	    refunded_amount = reserved_amount,
	    updated_at = NOW()
	WHERE id = $1
	  AND tenant_id = $2
	  AND quota_type = $3
	  AND status = 'pending'
	`, reservationID, tenantID, aiPointsQuotaType)
	if err != nil {
		return false, err
	}
	if tag.RowsAffected() > 0 {
		return true, nil
	}
	return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID)
}

func ignoreAIPointReservationReset(ctx context.Context, tx pgx.Tx, updated bool, tenantID, reservationID int64) error {
	if updated {
		return nil
	}
	var status string
	err := tx.QueryRow(ctx, `
	SELECT status
	FROM quota_reservations
	WHERE id = $1
	  AND tenant_id = $2
	  AND quota_type = $3
	`, reservationID, tenantID, aiPointsQuotaType).Scan(&status)
	if err != nil {
		return err
	}
	if status == "reset" {
		return nil
	}
	return errors.New("ai points reservation is not pending")
}