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

package repository

import (
	"context"
	"time"

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

type KolUsageLog struct {
	ID               int64
	TenantID         int64
	GenerationTaskID *int64
	Status           string
	CreatedAt        time.Time
}

type KolUsageByTaskRecord struct {
	ID        int64
	TenantID  int64
	PackageID int64
	PromptID  int64
	Status    string
	CreatedAt time.Time
}

type CreateKolUsageLogInput struct {
	TenantID             int64
	SubscriptionID       *int64
	SubscriptionPromptID *int64
	PackageID            int64
	PromptID             int64
	PromptRevisionNo     int
	GenerationTaskID     *int64
	VariablesSnapshot    []byte
	SchemaSnapshot       []byte
	RenderedPromptHash   *string
}

type KolDashboardOverview struct {
	TotalSubs    int64
	MonthSubs    int64
	TotalUsage   int64
	MonthUsage   int64
	FailureCount int64
}

type KolDashboardTrendPoint struct {
	Day           time.Time
	Usages        int64
	Subscriptions int64
}

type KolUsagePackageStats struct {
	PackageID int64
	Total     int64
	Completed int64
	Failed    int64
}

type KolUsageRepository interface {
	Create(ctx context.Context, input CreateKolUsageLogInput) (*KolUsageLog, error)
	MarkCompletedByTask(ctx context.Context, tenantID, taskID, articleID int64) error
	MarkFailedByTask(ctx context.Context, tenantID, taskID int64, errorMsg string) error
	GetByGenerationTask(ctx context.Context, tenantID, taskID int64) (*KolUsageByTaskRecord, error)
	DashboardOverview(ctx context.Context, packageIDs []int64) (*KolDashboardOverview, error)
	DashboardTrend(ctx context.Context, packageIDs []int64, periodDays int) ([]KolDashboardTrendPoint, error)
	PerPackageStats(ctx context.Context, packageIDs []int64) ([]KolUsagePackageStats, error)
}

type kolUsageRepository struct {
	q generated.Querier
}

func NewKolUsageRepository(db generated.DBTX) KolUsageRepository {
	return &kolUsageRepository{q: newQuerier(db)}
}

func (r *kolUsageRepository) Create(ctx context.Context, input CreateKolUsageLogInput) (*KolUsageLog, error) {
	row, err := r.q.CreateKolUsageLog(ctx, generated.CreateKolUsageLogParams{
		TenantID:             input.TenantID,
		SubscriptionID:       pgInt8(input.SubscriptionID),
		SubscriptionPromptID: pgInt8(input.SubscriptionPromptID),
		PackageID:            input.PackageID,
		PromptID:             input.PromptID,
		PromptRevisionNo:     int32(input.PromptRevisionNo),
		GenerationTaskID:     pgInt8(input.GenerationTaskID),
		VariablesSnapshot:    input.VariablesSnapshot,
		SchemaSnapshot:       input.SchemaSnapshot,
		RenderedPromptHash:   pgText(input.RenderedPromptHash),
	})
	if err != nil {
		return nil, err
	}
	return &KolUsageLog{
		ID:               row.ID,
		TenantID:         row.TenantID,
		GenerationTaskID: nullableInt64(row.GenerationTaskID),
		Status:           row.Status,
		CreatedAt:        timeFromTimestamp(row.CreatedAt),
	}, nil
}

func (r *kolUsageRepository) MarkCompletedByTask(ctx context.Context, tenantID, taskID, articleID int64) error {
	return r.q.MarkKolUsageCompletedByTask(ctx, generated.MarkKolUsageCompletedByTaskParams{
		ArticleID:        pgInt8(&articleID),
		GenerationTaskID: pgInt8(&taskID),
		TenantID:         tenantID,
	})
}

func (r *kolUsageRepository) MarkFailedByTask(ctx context.Context, tenantID, taskID int64, errorMsg string) error {
	return r.q.MarkKolUsageFailedByTask(ctx, generated.MarkKolUsageFailedByTaskParams{
		ErrorMessage:     pgText(&errorMsg),
		GenerationTaskID: pgInt8(&taskID),
		TenantID:         tenantID,
	})
}

func (r *kolUsageRepository) GetByGenerationTask(ctx context.Context, tenantID, taskID int64) (*KolUsageByTaskRecord, error) {
	row, err := r.q.GetKolUsageByGenerationTask(ctx, generated.GetKolUsageByGenerationTaskParams{
		GenerationTaskID: pgInt8(&taskID),
		TenantID:         tenantID,
	})
	if err != nil {
		// Callers use errors.Is(err, pgx.ErrNoRows) for "not found"; preserve err semantics.
		return nil, err
	}
	return &KolUsageByTaskRecord{
		ID:        row.ID,
		TenantID:  row.TenantID,
		PackageID: row.PackageID,
		PromptID:  row.PromptID,
		Status:    row.Status,
		CreatedAt: timeFromTimestamp(row.CreatedAt),
	}, nil
}

func (r *kolUsageRepository) DashboardOverview(ctx context.Context, packageIDs []int64) (*KolDashboardOverview, error) {
	if len(packageIDs) == 0 {
		return &KolDashboardOverview{}, nil
	}
	row, err := r.q.KolDashboardOverview(ctx, packageIDs)
	if err != nil {
		return nil, err
	}
	return &KolDashboardOverview{
		TotalSubs:    row.TotalSubs,
		MonthSubs:    row.MonthSubs,
		TotalUsage:   row.TotalUsage,
		MonthUsage:   row.MonthUsage,
		FailureCount: row.FailureCount,
	}, nil
}

func (r *kolUsageRepository) DashboardTrend(ctx context.Context, packageIDs []int64, periodDays int) ([]KolDashboardTrendPoint, error) {
	if len(packageIDs) == 0 {
		return []KolDashboardTrendPoint{}, nil
	}
	rows, err := r.q.KolDashboardTrend(ctx, generated.KolDashboardTrendParams{
		PackageIds: packageIDs,
		PeriodDays: int32(periodDays),
	})
	if err != nil {
		return nil, err
	}
	result := make([]KolDashboardTrendPoint, len(rows))
	for i, row := range rows {
		var day time.Time
		if row.D.Valid {
			day = row.D.Time
		}
		result[i] = KolDashboardTrendPoint{
			Day:           day,
			Usages:        row.Usages,
			Subscriptions: row.Subscriptions,
		}
	}
	return result, nil
}

func (r *kolUsageRepository) PerPackageStats(ctx context.Context, packageIDs []int64) ([]KolUsagePackageStats, error) {
	if len(packageIDs) == 0 {
		return []KolUsagePackageStats{}, nil
	}
	rows, err := r.q.KolUsageCountByPackage(ctx, packageIDs)
	if err != nil {
		return nil, err
	}
	result := make([]KolUsagePackageStats, len(rows))
	for i, row := range rows {
		result[i] = KolUsagePackageStats{
			PackageID: row.PackageID,
			Total:     row.Total,
			Completed: row.Completed,
			Failed:    row.Failed,
		}
	}
	return result, nil
}