geo/server/internal/scheduler/ai_point_usage_cleanup_worker.go

package scheduler

import (
	"context"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	"go.uber.org/zap"

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

const (
	defaultAIPointUsageCleanupBatch   = 5000
	defaultAIPointUsageCleanupBatches = 100
)

type AIPointUsageCleanupWorker struct {
	pool   *pgxpool.Pool
	logger *zap.Logger
}

func NewAIPointUsageCleanupWorker(pool *pgxpool.Pool, logger *zap.Logger) *AIPointUsageCleanupWorker {
	return &AIPointUsageCleanupWorker{pool: pool, logger: logger}
}

func (w *AIPointUsageCleanupWorker) RunOnce(ctx context.Context, run JobRunContext) (map[string]any, error) {
	if w == nil || w.pool == nil {
		return map[string]any{"skipped": true, "reason": "worker_not_ready"}, nil
	}

	startedAt := time.Now()
	batchSize := AsInt(run.Config, "batch_size", defaultAIPointUsageCleanupBatch)
	if run.Job != nil && run.Job.BatchSize != nil {
		batchSize = *run.Job.BatchSize
	}
	if batchSize <= 0 {
		batchSize = defaultAIPointUsageCleanupBatch
	}
	maxBatches := AsInt(run.Config, "max_batches_per_run", defaultAIPointUsageCleanupBatches)
	if maxBatches <= 0 {
		maxBatches = defaultAIPointUsageCleanupBatches
	}
	statementTimeout := AsDuration(run.Config, "statement_timeout", 5*time.Second)
	dryRun := run.DryRun
	now := time.Now().UTC()

	stats := map[string]any{
		"batch_size":          batchSize,
		"max_batches_per_run": maxBatches,
		"dry_run":             dryRun,
	}

	conn, err := w.pool.Acquire(ctx)
	if err != nil {
		return stats, fmt.Errorf("acquire tenant connection: %w", err)
	}
	defer conn.Release()

	if statementTimeout > 0 {
		if _, err := conn.Exec(ctx, `SELECT set_config('statement_timeout', $1, false)`, fmt.Sprintf("%dms", statementTimeout.Milliseconds())); err != nil {
			return stats, fmt.Errorf("set statement timeout: %w", err)
		}
		defer conn.Exec(context.Background(), `RESET statement_timeout`)
	}

	tenantIDs, err := w.loadTenantIDs(ctx, conn)
	if err != nil {
		return stats, fmt.Errorf("load ai point usage tenants: %w", err)
	}

	planRepo := tenantrepo.NewTenantPlanRepository(conn)
	totalDeleted := int64(0)
	totalCandidateProbe := int64(0)
	tenantCount := 0
	cleanedTenantCount := 0
	skippedNoPlan := 0
	batches := 0
	batchLimitReached := false

	for _, tenantID := range tenantIDs {
		if ctx.Err() != nil {
			return stats, ErrorFromContext(ctx)
		}
		if !dryRun && batches >= maxBatches {
			batchLimitReached = true
			break
		}

		tenantCount++
		access, err := planRepo.GetTenantPlanAccess(ctx, tenantID, now)
		if err != nil {
			return stats, fmt.Errorf("load tenant %d plan access: %w", tenantID, err)
		}
		cutoff, ok := aiPointUsageCleanupCutoff(access, now)
		if !ok {
			skippedNoPlan++
			continue
		}

		if dryRun {
			candidates, err := w.probeTenant(ctx, conn, tenantID, cutoff, batchSize)
			if err != nil {
				return stats, fmt.Errorf("probe tenant %d ai point usage: %w", tenantID, err)
			}
			if candidates > 0 {
				cleanedTenantCount++
				totalCandidateProbe += candidates
			}
			continue
		}

		deleted, tenantBatches, err := w.deleteTenant(ctx, conn, tenantID, cutoff, batchSize, maxBatches-batches)
		if err != nil {
			return stats, fmt.Errorf("delete tenant %d ai point usage: %w", tenantID, err)
		}
		if deleted > 0 {
			cleanedTenantCount++
			totalDeleted += deleted
			batches += tenantBatches
		}
	}

	stats["tenant_count"] = tenantCount
	stats["cleaned_tenant_count"] = cleanedTenantCount
	stats["skipped_no_plan_count"] = skippedNoPlan
	stats["deleted_count"] = totalDeleted
	stats["batch_count"] = batches
	stats["batch_limit_reached"] = batchLimitReached
	if dryRun {
		stats["candidate_probe_count"] = totalCandidateProbe
	}
	stats["duration_ms"] = time.Since(startedAt).Milliseconds()
	return stats, nil
}

func (w *AIPointUsageCleanupWorker) loadTenantIDs(ctx context.Context, conn *pgxpool.Conn) ([]int64, error) {
	rows, err := conn.Query(ctx, `
SELECT DISTINCT tenant_id
FROM ai_point_usage_logs
ORDER BY tenant_id
`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tenantIDs []int64
	for rows.Next() {
		var tenantID int64
		if err := rows.Scan(&tenantID); err != nil {
			return nil, err
		}
		tenantIDs = append(tenantIDs, tenantID)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return tenantIDs, nil
}

func (w *AIPointUsageCleanupWorker) probeTenant(ctx context.Context, conn *pgxpool.Conn, tenantID int64, cutoff time.Time, limit int) (int64, error) {
	if limit <= 0 {
		limit = defaultAIPointUsageCleanupBatch
	}
	var count int64
	if err := conn.QueryRow(ctx, aiPointUsageCleanupProbeSQL(), tenantID, cutoff.UTC(), limit).Scan(&count); err != nil {
		return 0, err
	}
	return count, nil
}

func (w *AIPointUsageCleanupWorker) deleteTenant(ctx context.Context, conn *pgxpool.Conn, tenantID int64, cutoff time.Time, batchSize, maxBatches int) (int64, int, error) {
	total := int64(0)
	batches := 0
	for batches < maxBatches {
		tag, err := conn.Exec(ctx, aiPointUsageCleanupDeleteSQL(), tenantID, cutoff.UTC(), batchSize)
		if err != nil {
			return total, batches, err
		}
		affected := tag.RowsAffected()
		if affected == 0 {
			break
		}
		total += affected
		batches++
		if affected < int64(batchSize) {
			break
		}
	}
	return total, batches, nil
}

func aiPointUsageCleanupCutoff(access *tenantrepo.TenantPlanAccess, now time.Time) (time.Time, bool) {
	if access == nil {
		return time.Time{}, false
	}

	referenceTime := now.UTC()
	if !access.EndAt.IsZero() && !access.EndAt.After(referenceTime) && access.EndAt.After(access.StartAt) {
		referenceTime = access.EndAt.Add(-time.Nanosecond)
	}
	if referenceTime.Before(access.StartAt) {
		referenceTime = access.StartAt
	}

	windowStart, _, _ := access.AIPointsWindow(referenceTime)
	if windowStart.IsZero() {
		return time.Time{}, false
	}
	return windowStart.UTC(), true
}

func aiPointUsageCleanupProbeSQL() string {
	return `
SELECT COUNT(*)::bigint
FROM (
	SELECT id
	FROM ai_point_usage_logs
	WHERE tenant_id = $1
	  AND created_at < $2
	  AND status IN ('completed', 'refunded', 'failed')
	ORDER BY created_at ASC, id ASC
	LIMIT $3
) candidates
`
}

func aiPointUsageCleanupDeleteSQL() string {
	return `
WITH doomed AS (
	SELECT id
	FROM ai_point_usage_logs
	WHERE tenant_id = $1
	  AND created_at < $2
	  AND status IN ('completed', 'refunded', 'failed')
	ORDER BY created_at ASC, id ASC
	LIMIT $3
	FOR UPDATE SKIP LOCKED
)
DELETE FROM ai_point_usage_logs l
USING doomed
WHERE l.id = doomed.id
`
}