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
}

// IncrementDailyStats increments the daily pull or push count for a repository
// Creates a new daily record if none exists for the current date, updates existing otherwise
// On first daily record for a repo, seeds it with existing cumulative stats so trend charts
// have a starting data point for pre-daily-tracking history
func (p *HoldPDS) IncrementDailyStats(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)
	}

	date := time.Now().UTC().Format("2006-01-02")
	rkey := atproto.DailyStatsRecordKey(ownerDID, repository, date)
	now := time.Now().Format(time.RFC3339)

	// Try to get existing daily record for today
	_, existing, err := p.GetDailyStats(ctx, ownerDID, repository, date)
	if err != nil {
		// No daily record for today — check if we need to seed from cumulative stats
		seedPull, seedPush := p.getSeedCounts(ctx, ownerDID, repository)

		record := atproto.NewDailyStatsRecord(ownerDID, repository, date)
		if operation == "pull" {
			record.PullCount = 1
		} else {
			record.PushCount = 1
		}
		record.UpdatedAt = now

		// If there are existing cumulative counts but no daily records yet,
		// create a seed record for the previous day with the historical totals
		if seedPull > 0 || seedPush > 0 {
			if err := p.seedDailyStats(ctx, ownerDID, repository, date, seedPull, seedPush); err != nil {
				slog.Warn("Failed to seed daily stats from cumulative",
					"ownerDID", ownerDID,
					"repository", repository,
					"error", err)
			}
		}

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

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

	// Record exists — increment
	if operation == "pull" {
		existing.PullCount++
	} else {
		existing.PushCount++
	}
	existing.UpdatedAt = now

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

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

// getSeedCounts returns the cumulative pull/push counts from io.atcr.hold.stats
// minus any already-tracked daily counts. Returns (0, 0) if no seeding is needed.
func (p *HoldPDS) getSeedCounts(ctx context.Context, ownerDID, repository string) (int64, int64) {
	// Check if any daily records already exist for this repo
	dailyStats, err := p.ListDailyStatsForRepo(ctx, ownerDID, repository)
	if err == nil && len(dailyStats) > 0 {
		// Daily records already exist — no seeding needed
		return 0, 0
	}

	// Get cumulative stats
	_, cumulative, err := p.GetStats(ctx, ownerDID, repository)
	if err != nil || cumulative == nil {
		return 0, 0
	}

	return cumulative.PullCount, cumulative.PushCount
}

// seedDailyStats creates a seed daily record with historical cumulative totals
// dated to the day before the first real daily record
func (p *HoldPDS) seedDailyStats(ctx context.Context, ownerDID, repository, currentDate string, pullCount, pushCount int64) error {
	// Parse current date and go back one day for the seed record
	t, err := time.Parse("2006-01-02", currentDate)
	if err != nil {
		return fmt.Errorf("failed to parse date: %w", err)
	}
	seedDate := t.AddDate(0, 0, -1).Format("2006-01-02")
	rkey := atproto.DailyStatsRecordKey(ownerDID, repository, seedDate)

	record := atproto.NewDailyStatsRecord(ownerDID, repository, seedDate)
	record.PullCount = pullCount
	record.PushCount = pushCount
	record.UpdatedAt = time.Now().Format(time.RFC3339)

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

	slog.Info("Seeded daily stats from cumulative totals",
		"ownerDID", ownerDID,
		"repository", repository,
		"seedDate", seedDate,
		"pullCount", pullCount,
		"pushCount", pushCount)
	return nil
}

// GetDailyStats retrieves the daily stats record for a repository on a specific date
func (p *HoldPDS) GetDailyStats(ctx context.Context, ownerDID, repository, date string) (cid.Cid, *atproto.DailyStatsRecord, error) {
	rkey := atproto.DailyStatsRecordKey(ownerDID, repository, date)

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

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

	return recordCID, dailyRecord, nil
}

// ListDailyStatsForRepo returns all daily stats records for a specific owner+repo
func (p *HoldPDS) ListDailyStatsForRepo(ctx context.Context, ownerDID, repository string) ([]*atproto.DailyStatsRecord, error) {
	session, err := p.carstore.ReadOnlySession(p.uid)
	if err != nil {
		return nil, fmt.Errorf("failed to get read-only session: %w", err)
	}

	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() {
		return []*atproto.DailyStatsRecord{}, nil
	}

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

	var stats []*atproto.DailyStatsRecord

	err = r.ForEach(ctx, atproto.DailyStatsCollection, func(k string, v cid.Cid) error {
		parts := strings.Split(k, "/")
		if len(parts) < 2 {
			return nil
		}

		actualCollection := strings.Join(parts[:len(parts)-1], "/")
		if actualCollection != atproto.DailyStatsCollection {
			return repo.ErrDoneIterating
		}

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

		if recBytes == nil {
			return nil
		}

		var dailyRecord atproto.DailyStatsRecord
		if err := dailyRecord.UnmarshalCBOR(bytes.NewReader(*recBytes)); err != nil {
			slog.Warn("Failed to unmarshal daily stats record", "key", k, "error", err)
			return nil
		}

		if dailyRecord.OwnerDID == ownerDID && dailyRecord.Repository == repository {
			stats = append(stats, &dailyRecord)
		}
		return nil
	})

	if err != nil {
		if errors.Is(err, repo.ErrDoneIterating) {
			// Expected
		} else if strings.Contains(err.Error(), "not found") {
			return []*atproto.DailyStatsRecord{}, nil
		} else {
			return nil, fmt.Errorf("failed to iterate daily stats records: %w", err)
		}
	}

	return stats, 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
}

// 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
}