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

package repository

import (
	"context"
	"time"

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

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

type AIPointUsageLog struct {
	ID                 int64
	TenantID           int64
	OperatorID         *int64
	QuotaReservationID *int64
	UsageType          string
	ResourceType       *string
	ResourceID         *int64
	ResourceUID        *string
	RequestChars       int
	BaseChars          int
	Points             int
	Status             string
	Model              *string
	MetadataJSON       []byte
	ErrorMessage       *string
	CompletedAt        *time.Time
	CreatedAt          time.Time
	UpdatedAt          time.Time
}

type CreateAIPointUsageLogInput struct {
	TenantID           int64
	OperatorID         int64
	QuotaReservationID int64
	UsageType          string
	ResourceType       *string
	ResourceID         *int64
	ResourceUID        *string
	RequestChars       int
	BaseChars          int
	Points             int
	MetadataJSON       []byte
}

type AIPointUsageListParams struct {
	TenantID int64
	Limit    int
	Offset   int
}

type AIPointUsageRepository interface {
	Create(ctx context.Context, input CreateAIPointUsageLogInput) (int64, error)
	MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error
	MarkRefunded(ctx context.Context, tenantID, id int64, errorMessage string) error
	FindPendingByResource(ctx context.Context, tenantID int64, resourceType string, resourceID *int64, resourceUID *string) (*AIPointUsageLog, error)
	List(ctx context.Context, params AIPointUsageListParams) ([]AIPointUsageLog, error)
	Count(ctx context.Context, tenantID int64) (int64, error)
}

type aiPointUsageRepository struct {
	db generated.DBTX
}

func NewAIPointUsageRepository(db generated.DBTX) AIPointUsageRepository {
	return &aiPointUsageRepository{db: db}
}

func (r *aiPointUsageRepository) Create(ctx context.Context, input CreateAIPointUsageLogInput) (int64, error) {
	metadataJSON := input.MetadataJSON
	if len(metadataJSON) == 0 {
		metadataJSON = []byte("{}")
	}
	var operatorID *int64
	if input.OperatorID > 0 {
		operatorID = &input.OperatorID
	}
	var reservationID *int64
	if input.QuotaReservationID > 0 {
		reservationID = &input.QuotaReservationID
	}

	var id int64
	err := r.db.QueryRow(ctx, `
INSERT INTO ai_point_usage_logs (
	tenant_id, operator_id, quota_reservation_id, usage_type,
	resource_type, resource_id, resource_uid,
	request_chars, base_chars, points, metadata_json
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::jsonb)
RETURNING id
`, input.TenantID, operatorID, reservationID, input.UsageType,
		input.ResourceType, input.ResourceID, input.ResourceUID,
		input.RequestChars, input.BaseChars, input.Points, metadataJSON,
	).Scan(&id)
	return id, err
}

func (r *aiPointUsageRepository) MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error {
	_, err := r.db.Exec(ctx, `
UPDATE ai_point_usage_logs
SET status = 'completed',
    model = $1,
    error_message = NULL,
    completed_at = COALESCE(completed_at, NOW()),
    updated_at = NOW()
WHERE id = $2
  AND tenant_id = $3
  AND status = 'pending'
`, model, id, tenantID)
	return err
}

func (r *aiPointUsageRepository) MarkRefunded(ctx context.Context, tenantID, id int64, errorMessage string) error {
	_, err := r.db.Exec(ctx, `
UPDATE ai_point_usage_logs
SET status = 'refunded',
    error_message = $1,
    completed_at = COALESCE(completed_at, NOW()),
    updated_at = NOW()
WHERE id = $2
  AND tenant_id = $3
  AND status = 'pending'
`, errorMessage, id, tenantID)
	return err
}

func (r *aiPointUsageRepository) FindPendingByResource(
	ctx context.Context,
	tenantID int64,
	resourceType string,
	resourceID *int64,
	resourceUID *string,
) (*AIPointUsageLog, error) {
	row := r.db.QueryRow(ctx, `
SELECT id, tenant_id, operator_id, quota_reservation_id, usage_type,
       resource_type, resource_id, resource_uid, request_chars, base_chars,
       points, status, model, metadata_json, error_message, completed_at,
       created_at, updated_at
FROM ai_point_usage_logs
WHERE tenant_id = $1
  AND resource_type = $2
  AND ($3::bigint IS NULL OR resource_id = $3)
  AND ($4::varchar IS NULL OR resource_uid = $4)
  AND status = 'pending'
ORDER BY created_at DESC, id DESC
LIMIT 1
`, tenantID, resourceType, resourceID, resourceUID)

	var item AIPointUsageLog
	err := scanAIPointUsageLog(row, &item)
	if err != nil {
		return nil, err
	}
	return &item, nil
}

func (r *aiPointUsageRepository) List(ctx context.Context, params AIPointUsageListParams) ([]AIPointUsageLog, error) {
	limit := params.Limit
	if limit <= 0 || limit > 100 {
		limit = 50
	}
	offset := params.Offset
	if offset < 0 {
		offset = 0
	}

	rows, err := r.db.Query(ctx, `
SELECT id, tenant_id, operator_id, quota_reservation_id, usage_type,
       resource_type, resource_id, resource_uid, request_chars, base_chars,
       points, status, model, metadata_json, error_message, completed_at,
       created_at, updated_at
FROM ai_point_usage_logs
WHERE tenant_id = $1
ORDER BY created_at DESC, id DESC
LIMIT $2 OFFSET $3
`, params.TenantID, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	items := make([]AIPointUsageLog, 0, limit)
	for rows.Next() {
		var item AIPointUsageLog
		if err := scanAIPointUsageLog(rows, &item); err != nil {
			return nil, err
		}
		items = append(items, item)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

func (r *aiPointUsageRepository) Count(ctx context.Context, tenantID int64) (int64, error) {
	var total int64
	err := r.db.QueryRow(ctx, `
	SELECT COUNT(*)
	FROM ai_point_usage_logs
	WHERE tenant_id = $1
	`, tenantID).Scan(&total)
	return total, err
}

type aiPointUsageScanner interface {
	Scan(dest ...any) error
}

func scanAIPointUsageLog(scanner aiPointUsageScanner, item *AIPointUsageLog) error {
	var operatorID pgtype.Int8
	var quotaReservationID pgtype.Int8
	var resourceType pgtype.Text
	var resourceID pgtype.Int8
	var resourceUID pgtype.Text
	var model pgtype.Text
	var errorMessage pgtype.Text
	var completedAt pgtype.Timestamptz

	if err := scanner.Scan(
		&item.ID,
		&item.TenantID,
		&operatorID,
		&quotaReservationID,
		&item.UsageType,
		&resourceType,
		&resourceID,
		&resourceUID,
		&item.RequestChars,
		&item.BaseChars,
		&item.Points,
		&item.Status,
		&model,
		&item.MetadataJSON,
		&errorMessage,
		&completedAt,
		&item.CreatedAt,
		&item.UpdatedAt,
	); err != nil {
		return err
	}

	item.OperatorID = nullableInt64(operatorID)
	item.QuotaReservationID = nullableInt64(quotaReservationID)
	item.ResourceType = nullableText(resourceType)
	item.ResourceID = nullableInt64(resourceID)
	item.ResourceUID = nullableText(resourceUID)
	item.Model = nullableText(model)
	item.ErrorMessage = nullableText(errorMessage)
	item.CompletedAt = optionalTime(completedAt)
	return nil
}