at-container-registry/pkg/hold/pds/stats.go

package pds

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"time"

	"atcr.io/pkg/atproto"
	"github.com/bluesky-social/indigo/repo"
	"github.com/ipfs/go-cid"
)

// IncrementStats increments the pull or push count for a repository
// operation should be "pull" or "push"
// Creates a new record if none exists, updates existing record otherwise
func (p *HoldPDS) IncrementStats(ctx context.Context, ownerDID, repository, operation string) error {
	if operation != "pull" && operation != "push" {
		return fmt.Errorf("invalid operation: %s (must be 'pull' or 'push')", operation)
	}

	rkey := atproto.StatsRecordKey(ownerDID, repository)
	now := time.Now().Format(time.RFC3339)

	// Try to get existing record
	_, existing, err := p.GetStats(ctx, ownerDID, repository)
	if err != nil {
		// Record doesn't exist - create new one
		record := atproto.NewStatsRecord(ownerDID, repository)
		if operation == "pull" {
			record.PullCount = 1
			record.LastPull = now
		} else {
			record.PushCount = 1
			record.LastPush = now
		}
		record.UpdatedAt = now

		_, _, err := p.repomgr.PutRecord(ctx, p.uid, atproto.StatsCollection, rkey, record)
		if err != nil {
			return fmt.Errorf("failed to create stats record: %w", err)
		}

		slog.Debug("Created stats record",
			"ownerDID", ownerDID,
			"repository", repository,
			"operation", operation)
		return nil
	}

	// Record exists - update it
	if operation == "pull" {
		existing.PullCount++
		existing.LastPull = now
	} else {
		existing.PushCount++
		existing.LastPush = now
	}
	existing.UpdatedAt = now

	_, err = p.repomgr.UpdateRecord(ctx, p.uid, atproto.StatsCollection, rkey, existing)
	if err != nil {
		return fmt.Errorf("failed to update stats record: %w", err)
	}

	slog.Debug("Updated stats record",
		"ownerDID", ownerDID,
		"repository", repository,
		"operation", operation,
		"pullCount", existing.PullCount,
		"pushCount", existing.PushCount)
	return nil
}

// GetStats retrieves the stats record for a repository
// Returns nil, nil if no stats record exists
func (p *HoldPDS) GetStats(ctx context.Context, ownerDID, repository string) (cid.Cid, *atproto.StatsRecord, error) {
	rkey := atproto.StatsRecordKey(ownerDID, repository)

	recordCID, val, err := p.repomgr.GetRecord(ctx, p.uid, atproto.StatsCollection, rkey, cid.Undef)
	if err != nil {
		return cid.Undef, nil, err
	}

	statsRecord, ok := val.(*atproto.StatsRecord)
	if !ok {
		return cid.Undef, nil, fmt.Errorf("unexpected type for stats record: %T", val)
	}

	return recordCID, statsRecord, nil
}

// SetStats directly sets the stats for a repository (used for migration)
// Creates or updates the stats record with the specified counts
func (p *HoldPDS) SetStats(ctx context.Context, ownerDID, repository string, pullCount, pushCount int64, lastPull, lastPush string) error {
	rkey := atproto.StatsRecordKey(ownerDID, repository)
	now := time.Now().Format(time.RFC3339)

	// Try to get existing record
	_, existing, err := p.GetStats(ctx, ownerDID, repository)
	if err != nil {
		// Record doesn't exist - create new one
		record := &atproto.StatsRecord{
			Type:       atproto.StatsCollection,
			OwnerDID:   ownerDID,
			Repository: repository,
			PullCount:  pullCount,
			PushCount:  pushCount,
			LastPull:   lastPull,
			LastPush:   lastPush,
			UpdatedAt:  now,
		}

		_, _, err := p.repomgr.PutRecord(ctx, p.uid, atproto.StatsCollection, rkey, record)
		if err != nil {
			return fmt.Errorf("failed to create stats record: %w", err)
		}
		return nil
	}

	// Record exists - update it
	existing.PullCount = pullCount
	existing.PushCount = pushCount
	existing.LastPull = lastPull
	existing.LastPush = lastPush
	existing.UpdatedAt = now

	_, err = p.repomgr.UpdateRecord(ctx, p.uid, atproto.StatsCollection, rkey, existing)
	if err != nil {
		return fmt.Errorf("failed to update stats record: %w", err)
	}

	return nil
}

// ListStats returns all stats records in the hold's PDS
// This is used by AppView to aggregate stats from all holds
func (p *HoldPDS) ListStats(ctx context.Context) ([]*atproto.StatsRecord, error) {
	// Get read-only session from carstore
	session, err := p.carstore.ReadOnlySession(p.uid)
	if err != nil {
		return nil, fmt.Errorf("failed to get read-only session: %w", err)
	}

	// Get repo head
	head, err := p.carstore.GetUserRepoHead(ctx, p.uid)
	if err != nil {
		return nil, fmt.Errorf("failed to get repo head: %w", err)
	}

	if !head.Defined() {
		// No repo yet, return empty list
		return []*atproto.StatsRecord{}, nil
	}

	// Open repo
	r, err := repo.OpenRepo(ctx, session, head)
	if err != nil {
		return nil, fmt.Errorf("failed to open repo: %w", err)
	}

	var stats []*atproto.StatsRecord

	// Iterate over all stats records
	err = r.ForEach(ctx, atproto.StatsCollection, func(k string, v cid.Cid) error {
		// Extract collection and rkey from full path (k is like "io.atcr.hold.stats/abcd1234...")
		parts := strings.Split(k, "/")
		if len(parts) < 2 {
			return nil // Skip invalid keys
		}

		// Extract actual collection
		actualCollection := strings.Join(parts[:len(parts)-1], "/")

		// MST keys are sorted, so once we hit a different collection, stop walking
		if actualCollection != atproto.StatsCollection {
			return repo.ErrDoneIterating
		}

		// Get record bytes
		_, recBytes, err := r.GetRecordBytes(ctx, k)
		if err != nil {
			slog.Warn("Failed to get stats record bytes", "key", k, "error", err)
			return nil // Continue with other records
		}

		if recBytes == nil {
			return nil
		}

		// Unmarshal the CBOR bytes
		var statsRecord atproto.StatsRecord
		if err := statsRecord.UnmarshalCBOR(bytes.NewReader(*recBytes)); err != nil {
			slog.Warn("Failed to unmarshal stats record", "key", k, "error", err)
			return nil // Continue with other records
		}

		stats = append(stats, &statsRecord)
		return nil
	})

	if err != nil {
		// ErrDoneIterating is expected when we stop walking early
		if errors.Is(err, repo.ErrDoneIterating) {
			// Successfully stopped at collection boundary
		} else if strings.Contains(err.Error(), "not found") {
			// Collection doesn't exist yet - return empty list
			return []*atproto.StatsRecord{}, nil
		} else {
			return nil, fmt.Errorf("failed to iterate stats records: %w", err)
		}
	}

	return stats, nil
}

// ListStatsRecordsForUser returns all stats records where the user is the repository owner
// Used for GDPR data export to return all stats for repositories owned by the user
func (p *HoldPDS) ListStatsRecordsForUser(ctx context.Context, userDID string) ([]*atproto.StatsRecord, error) {
	// Get all stats records and filter by ownerDID
	allStats, err := p.ListStats(ctx)
	if err != nil {
		return nil, err
	}

	var userStats []*atproto.StatsRecord
	for _, stat := range allStats {
		if stat.OwnerDID == userDID {
			userStats = append(userStats, stat)
		}
	}

	if userStats == nil {
		userStats = []*atproto.StatsRecord{}
	}

	return userStats, nil
}