pool-stats/database/db.go

package database

import (
	"fmt"
	"log"
	"pool-stats/helpers"
	"pool-stats/models"
	"sort"
	"time"

	"github.com/ostafen/clover/v2"
	"github.com/ostafen/clover/v2/document"
	c "github.com/ostafen/clover/v2/query"
	badgerstore "github.com/ostafen/clover/v2/store/badger"
)

const (
	CollectionName                    = "shares"
	TopSharesCollectionName           = "TopShares"
	TimeWindowHighShareCollectionName = "TimeWindowHighShareStat"
	DailyStatsCollectionName          = "DailyStats"
)

func InitDatabase(path string) (*clover.DB, error) {
	store, err := badgerstore.Open(path)
	if err != nil {
		return nil, fmt.Errorf("failed to open BadgerDB store: %v", err)
	}

	db, err := clover.OpenWithStore(store)
	if err != nil {
		return nil, fmt.Errorf("failed to open CloverDB: %v", err)
	}

	// Ensure collection exists
	hasCollection, err := db.HasCollection(CollectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to check collection: %v", err)
	}

	if !hasCollection {
		if err := db.CreateCollection(CollectionName); err != nil {
			return nil, fmt.Errorf("failed to create collection: %v", err)
		}
		if err := db.CreateIndex(CollectionName, "CreateDate"); err != nil {
			return nil, fmt.Errorf("failed to create index: %v", err)
		}
	}

	// Init TopShares collection
	hasTopSharesCollection, err := db.HasCollection(TopSharesCollectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to check TopShares collection: %v", err)
	}

	if !hasTopSharesCollection {
		if err := db.CreateCollection(TopSharesCollectionName); err != nil {
			return nil, fmt.Errorf("failed to create TopShares collection: %v", err)
		}

		if err := db.CreateIndex(TopSharesCollectionName, "CreateDate"); err != nil {
			return nil, fmt.Errorf("failed to create index for TopShares: %v", err)
		}

		if err := db.CreateIndex(TopSharesCollectionName, "SDiff"); err != nil {
			return nil, fmt.Errorf("failed to create index for TopShares SDiff: %v", err)
		}
	}

	// Init TimeWindowHighShareStat collection
	hasTimeWindowCollection, err := db.HasCollection(TimeWindowHighShareCollectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to check TimeWindowHighShare collection: %v", err)
	}

	if !hasTimeWindowCollection {
		if err := db.CreateCollection(TimeWindowHighShareCollectionName); err != nil {
			return nil, fmt.Errorf("failed to create TimeWindowHighShare collection: %v", err)
		}
		if err := db.CreateIndex(TimeWindowHighShareCollectionName, "TimeWindowID"); err != nil {
			return nil, fmt.Errorf("failed to create index for TimeWindowHighShare: %v", err)
		}
	}

	// Init DailyStats collection
	hasDailyStatsCollection, err := db.HasCollection(DailyStatsCollectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to check DailyStats collection: %v", err)
	}
	if !hasDailyStatsCollection {
		if err := db.CreateCollection(DailyStatsCollectionName); err != nil {
			return nil, fmt.Errorf("failed to create DailyStats collection: %v", err)
		}
		if err := db.CreateIndex(DailyStatsCollectionName, "Date"); err != nil {
			return nil, fmt.Errorf("failed to create index for DailyStats: %v", err)
		}
	}

	return db, nil
}

func GetHighestSharesInRange(db *clover.DB, collection string, since time.Time, count int) ([]models.ShareLog, error) {
	// Convert `since` to the format in `createdate`
	lower := since.Unix()
	upper := time.Now().Unix()

	// Filter by timestamp range
	criteria := c.Field("CreateDate").GtEq(fmt.Sprint(lower)).
		And(c.Field("CreateDate").LtEq(fmt.Sprint(upper)))

	// Query sorted by "sdiff" descending, limit 1
	results, err := db.FindAll(c.NewQuery(collection).
		Where(criteria).
		Sort(c.SortOption{Field: "SDiff", Direction: -1}).
		Limit(count))

	if err != nil || len(results) == 0 {
		return nil, err
	}

	var shares []models.ShareLog
	for _, doc := range results {
		var s models.ShareLog
		if err := doc.Unmarshal(&s); err != nil {
			return nil, err
		}
		shares = append(shares, s)
	}

	return shares, nil
}

func PrintAllHashes(db *clover.DB) {
	docs, err := db.FindAll(c.NewQuery(CollectionName))
	if err != nil {
		log.Fatalf("Failed to read from collection: %v", err)
	}

	for _, doc := range docs {
		hash := doc.Get("Hash")
		fmt.Println(hash)
	}
}

func ListShares(db *clover.DB, offset int, count int) []models.ShareLog {
	results, err := db.FindAll(
		c.NewQuery(CollectionName).
			Sort(c.SortOption{Field: "CreateDate", Direction: -1}).
			Skip(offset).
			Limit(count),
	)
	if err != nil {
		log.Printf("failed to list shares: %v", err)
		return nil
	}

	shareLogs := make([]models.ShareLog, len(results))
	for idx, doc := range results {
		var shareLog models.ShareLog
		doc.Unmarshal(&shareLog)
		shareLogs[idx] = shareLog
	}

	return shareLogs
}

func ListTopShares(db *clover.DB) []models.ShareLog {
	results, err := db.FindAll(
		c.NewQuery(TopSharesCollectionName).
			Sort(c.SortOption{Field: "SDiff", Direction: -1}),
	)
	if err != nil {
		log.Printf("failed to list top shares: %v", err)
		return nil
	}

	topShares := make([]models.ShareLog, len(results))
	for idx, doc := range results {
		var shareLog models.ShareLog
		doc.Unmarshal(&shareLog)
		topShares[idx] = shareLog
	}

	return topShares
}

func ReplaceTopShares(db *clover.DB, shares []models.ShareLog) {
	db.Delete(c.NewQuery(TopSharesCollectionName))

	for _, share := range shares {
		doc := document.NewDocumentOf(&share)
		if _, err := db.InsertOne(TopSharesCollectionName, doc); err != nil {
			return
		}
	}
}

func GetTimeWindowHighShares(db *clover.DB) []models.TimeWindowHighShare {
	results, err := db.FindAll(
		c.NewQuery(TimeWindowHighShareCollectionName).
			Sort(c.SortOption{Field: "TimeWindowID", Direction: 1}),
	)
	if err != nil {
		log.Printf("failed to list time window high shares: %v", err)
		return nil
	}

	timeWindowHighShares := make([]models.TimeWindowHighShare, len(results))
	for idx, doc := range results {
		var timeWindowHighShare models.TimeWindowHighShare
		doc.Unmarshal(&timeWindowHighShare)
		timeWindowHighShares[idx] = timeWindowHighShare
	}

	return timeWindowHighShares
}

func SetTimeWindowHighShare(db *clover.DB, share models.TimeWindowHighShare) error {
	doc := document.NewDocumentOf(&share)

	existingDoc, _ := db.FindFirst(c.NewQuery(TimeWindowHighShareCollectionName).
		Where(c.Field("TimeWindowID").Eq(share.TimeWindowID)))
	if existingDoc != nil {
		db.ReplaceById(TimeWindowHighShareCollectionName, existingDoc.ObjectId(), doc)
	} else {
		db.InsertOne(TimeWindowHighShareCollectionName, doc)
	}

	return nil
}

func ListSharesInTimeRange(db *clover.DB, since time.Time, till time.Time) []models.ShareLog {
	lower := since.Unix()
	upper := till.Unix()

	results, err := db.FindAll(c.NewQuery(CollectionName).
		Where(c.Field("CreateDate").GtEq(fmt.Sprint(lower)).
			And(c.Field("CreateDate").LtEq(fmt.Sprint(upper)))).
		Sort(c.SortOption{Field: "CreateDate", Direction: -1}))

	if err != nil {
		log.Printf("failed to list shares in time range: %v", err)
		return nil
	}

	shareLogs := make([]models.ShareLog, len(results))
	for idx, doc := range results {
		var shareLog models.ShareLog
		doc.Unmarshal(&shareLog)
		shareLogs[idx] = shareLog
	}

	return shareLogs
}

// GetStatsForDay retrieves daily statistics for a given date
// Tries to find from DailyStats collection, calculates on the fly if not found and stores
func GetDailyStats(db *clover.DB, date time.Time) (*models.DailyStats, error) {
	dateStr := date.Format(time.DateOnly)

	// Check if stats already exist
	isToday := dateStr == time.Now().UTC().Format(time.DateOnly)
	existingDoc, err := db.FindFirst(c.NewQuery(DailyStatsCollectionName).
		Where(c.Field("Date").Eq(dateStr)))
	if !isToday && err == nil && existingDoc != nil {
		var stats models.DailyStats
		if err := existingDoc.Unmarshal(&stats); err != nil {
			return nil, fmt.Errorf("failed to unmarshal daily stats: %v", err)
		}
		return &stats, nil
	}

	// Get shares in range
	since := date.Truncate(24 * time.Hour)
	till := since.Add(24 * time.Hour)
	shares := ListSharesInTimeRange(db, since, till)
	sort.Slice(shares, func(i, j int) bool {
		return shares[i].SDiff > shares[j].SDiff
	})

	// Calculate daily stats
	stats := &models.DailyStats{
		Date:       dateStr,
		ShareCount: len(shares),
		Workers:    make(map[string]models.WorkerDailyStats),
	}

	if len(shares) > 0 {
		stats.TopShare = shares[0]
		stats.PoolHashrate = helpers.CalculateAverageHashrate(shares)
	}

	// Calculate worker stats
	sharesByWorker := make(map[string][]models.ShareLog)
	for _, share := range shares {
		sharesByWorker[share.WorkerName] = append(sharesByWorker[share.WorkerName], share)
	}
	for workerName, workerShares := range sharesByWorker {
		workerHashrate := helpers.CalculateAverageHashrate(workerShares)
		workerTopShare := workerShares[0] // Already sorted by SDiff

		stats.Workers[workerName] = models.WorkerDailyStats{
			TopShare: workerTopShare,
			Hashrate: workerHashrate,
			Shares:   len(workerShares),
		}
	}

	// Insert or update the daily stats in the collection
	doc := document.NewDocumentOf(stats)
	if _, err := db.InsertOne(DailyStatsCollectionName, doc); err != nil {
		return nil, fmt.Errorf("failed to insert daily stats: %v", err)
	}

	return stats, nil
}

func ClearDailyStats(db *clover.DB) error {
	// Delete all documents in DailyStats collection
	if err := db.Delete(c.NewQuery(DailyStatsCollectionName)); err != nil {
		return fmt.Errorf("failed to clear DailyStats collection: %v", err)
	}
	return nil
}