pool-stats/database/db.go

package database

import (
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"log"
	"pool-stats/helpers"
	"pool-stats/models"
	"sort"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

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

//go:embed schema.sql
var schemaSQL string

func InitDatabase(dsn string) (*pgxpool.Pool, error) {
	ctx := context.Background()

	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to PostgreSQL: %w", err)
	}

	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("failed to ping PostgreSQL: %w", err)
	}

	if _, err := pool.Exec(ctx, schemaSQL); err != nil {
		pool.Close()
		return nil, fmt.Errorf("failed to apply schema: %w", err)
	}

	return pool, nil
}

func shareColumns() string {
	return `work_info_id, client_id, enonce1, nonce2, nonce, ntime,
		diff, sdiff, hash, result, errn,
		create_date, create_ts, create_by, create_code, create_inet,
		worker_name, username, address, agent`
}

func scanShare(row pgx.Row) (models.ShareLog, error) {
	var s models.ShareLog
	var createTs time.Time
	err := row.Scan(
		&s.WorkInfoID, &s.ClientID, &s.Enonce1, &s.Nonce2, &s.Nonce, &s.NTime,
		&s.Diff, &s.SDiff, &s.Hash, &s.Result, &s.Errn,
		&s.CreateDate, &createTs, &s.CreateBy, &s.CreateCode, &s.CreateInet,
		&s.WorkerName, &s.Username, &s.Address, &s.Agent,
	)
	return s, err
}

func scanShares(rows pgx.Rows) ([]models.ShareLog, error) {
	defer rows.Close()
	var shares []models.ShareLog
	for rows.Next() {
		var s models.ShareLog
		var createTs time.Time
		if err := rows.Scan(
			&s.WorkInfoID, &s.ClientID, &s.Enonce1, &s.Nonce2, &s.Nonce, &s.NTime,
			&s.Diff, &s.SDiff, &s.Hash, &s.Result, &s.Errn,
			&s.CreateDate, &createTs, &s.CreateBy, &s.CreateCode, &s.CreateInet,
			&s.WorkerName, &s.Username, &s.Address, &s.Agent,
		); err != nil {
			return nil, err
		}
		shares = append(shares, s)
	}
	return shares, rows.Err()
}

func shareValues(s models.ShareLog) ([]any, error) {
	createTs, err := s.ParseCreateDate()
	if err != nil {
		createTs = helpers.ParseCreateDate(s.CreateDate)
	}
	if createTs.IsZero() {
		createTs = time.Unix(0, 0)
	}
	return []any{
		s.WorkInfoID, s.ClientID, s.Enonce1, s.Nonce2, s.Nonce, s.NTime,
		s.Diff, s.SDiff, s.Hash, s.Result, s.Errn,
		s.CreateDate, createTs, s.CreateBy, s.CreateCode, s.CreateInet,
		s.WorkerName, s.Username, s.Address, s.Agent,
	}, nil
}

func tableForCollection(collection string) string {
	switch collection {
	case TopSharesCollectionName:
		return "top_shares"
	default:
		return "shares"
	}
}

func InsertShare(db *pgxpool.Pool, share models.ShareLog) error {
	ctx := context.Background()
	vals, err := shareValues(share)
	if err != nil {
		return err
	}
	_, err = db.Exec(ctx, `
		INSERT INTO shares (`+shareColumns()+`)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20)`,
		vals...)
	return err
}

func GetHighestSharesInRange(db *pgxpool.Pool, collection string, since time.Time, count int) ([]models.ShareLog, error) {
	ctx := context.Background()
	table := tableForCollection(collection)
	rows, err := db.Query(ctx, `
		SELECT `+shareColumns()+`
		FROM `+table+`
		WHERE create_ts >= $1 AND create_ts <= $2
		ORDER BY sdiff DESC
		LIMIT $3`,
		since, time.Now(), count)
	if err != nil {
		return nil, err
	}
	return scanShares(rows)
}

func ListShares(db *pgxpool.Pool, offset int, count int) []models.ShareLog {
	ctx := context.Background()
	rows, err := db.Query(ctx, `
		SELECT `+shareColumns()+`
		FROM shares
		ORDER BY create_ts DESC
		OFFSET $1 LIMIT $2`,
		offset, count)
	if err != nil {
		log.Printf("failed to list shares: %v", err)
		return nil
	}
	shares, err := scanShares(rows)
	if err != nil {
		log.Printf("failed to scan shares: %v", err)
		return nil
	}
	return shares
}

func ListTopShares(db *pgxpool.Pool) []models.ShareLog {
	ctx := context.Background()
	rows, err := db.Query(ctx, `
		SELECT `+shareColumns()+`
		FROM top_shares
		ORDER BY sdiff DESC`)
	if err != nil {
		log.Printf("failed to list top shares: %v", err)
		return nil
	}
	shares, err := scanShares(rows)
	if err != nil {
		log.Printf("failed to scan top shares: %v", err)
		return nil
	}
	return shares
}

func ReplaceTopShares(db *pgxpool.Pool, shares []models.ShareLog) {
	ctx := context.Background()
	tx, err := db.Begin(ctx)
	if err != nil {
		log.Printf("failed to begin replace top shares: %v", err)
		return
	}
	defer tx.Rollback(ctx)

	if _, err := tx.Exec(ctx, `DELETE FROM top_shares`); err != nil {
		log.Printf("failed to clear top shares: %v", err)
		return
	}

	for _, share := range shares {
		vals, err := shareValues(share)
		if err != nil {
			log.Printf("failed to prepare top share: %v", err)
			return
		}
		if _, err := tx.Exec(ctx, `
			INSERT INTO top_shares (`+shareColumns()+`)
			VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20)`,
			vals...); err != nil {
			log.Printf("failed to insert top share: %v", err)
			return
		}
	}

	if err := tx.Commit(ctx); err != nil {
		log.Printf("failed to commit top shares: %v", err)
	}
}

func GetTimeWindowHighShares(db *pgxpool.Pool) []models.TimeWindowHighShare {
	ctx := context.Background()
	rows, err := db.Query(ctx, `
		SELECT time_window_id, time_window_name, sdiff, time
		FROM time_window_high_shares
		ORDER BY time_window_id ASC`)
	if err != nil {
		log.Printf("failed to list time window high shares: %v", err)
		return nil
	}
	defer rows.Close()

	var result []models.TimeWindowHighShare
	for rows.Next() {
		var tw models.TimeWindowHighShare
		if err := rows.Scan(&tw.TimeWindowID, &tw.TimeWindowName, &tw.SDiff, &tw.Time); err != nil {
			log.Printf("failed to scan time window high share: %v", err)
			return nil
		}
		result = append(result, tw)
	}
	return result
}

func SetTimeWindowHighShare(db *pgxpool.Pool, share models.TimeWindowHighShare) error {
	ctx := context.Background()
	_, err := db.Exec(ctx, `
		INSERT INTO time_window_high_shares (time_window_id, time_window_name, sdiff, time)
		VALUES ($1, $2, $3, $4)
		ON CONFLICT (time_window_id) DO UPDATE SET
			time_window_name = EXCLUDED.time_window_name,
			sdiff = EXCLUDED.sdiff,
			time = EXCLUDED.time`,
		share.TimeWindowID, share.TimeWindowName, share.SDiff, share.Time)
	return err
}

func ListSharesInTimeRange(db *pgxpool.Pool, since time.Time, till time.Time) []models.ShareLog {
	ctx := context.Background()
	rows, err := db.Query(ctx, `
		SELECT `+shareColumns()+`
		FROM shares
		WHERE create_ts >= $1 AND create_ts <= $2
		ORDER BY create_ts DESC`,
		since, till)
	if err != nil {
		log.Printf("failed to list shares in time range: %v", err)
		return nil
	}
	shares, err := scanShares(rows)
	if err != nil {
		log.Printf("failed to scan shares in time range: %v", err)
		return nil
	}
	return shares
}

func GetDailyStats(db *pgxpool.Pool, date time.Time) (*models.DailyStats, error) {
	ctx := context.Background()
	dateStr := date.Format(time.DateOnly)

	var (
		stats     models.DailyStats
		topShare  []byte
		workers   []byte
		expiresAt *time.Time
	)
	err := db.QueryRow(ctx, `
		SELECT date::text, share_count, top_share, pool_hashrate, workers, expires_at
		FROM daily_stats
		WHERE date = $1::date`, dateStr).Scan(
		&stats.Date, &stats.ShareCount, &topShare, &stats.PoolHashrate, &workers, &expiresAt,
	)
	if err == nil {
		validCache := expiresAt == nil || expiresAt.After(time.Now())
		if validCache {
			if err := json.Unmarshal(topShare, &stats.TopShare); err != nil {
				return nil, fmt.Errorf("failed to unmarshal top share: %w", err)
			}
			if err := json.Unmarshal(workers, &stats.Workers); err != nil {
				return nil, fmt.Errorf("failed to unmarshal workers: %w", err)
			}
			return &stats, nil
		}
		_ = DeleteDailyStatsForDay(db, date)
	} else if err != pgx.ErrNoRows {
		return nil, fmt.Errorf("failed to query daily stats: %w", err)
	}

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

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

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

	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)
		sort.Slice(workerShares, func(i, j int) bool {
			return workerShares[i].SDiff > workerShares[j].SDiff
		})
		computed.Workers[workerName] = models.WorkerDailyStats{
			TopShare: workerShares[0],
			Hashrate: workerHashrate,
			Shares:   len(workerShares),
		}
	}

	topShareJSON, err := json.Marshal(computed.TopShare)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal top share: %w", err)
	}
	workersJSON, err := json.Marshal(computed.Workers)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal workers: %w", err)
	}

	var expires any
	isToday := dateStr == time.Now().UTC().Format(time.DateOnly)
	if isToday {
		expires = time.Now().Add(5 * time.Minute)
	}

	_, err = db.Exec(ctx, `
		INSERT INTO daily_stats (date, share_count, top_share, pool_hashrate, workers, expires_at)
		VALUES ($1::date, $2, $3, $4, $5, $6)
		ON CONFLICT (date) DO UPDATE SET
			share_count = EXCLUDED.share_count,
			top_share = EXCLUDED.top_share,
			pool_hashrate = EXCLUDED.pool_hashrate,
			workers = EXCLUDED.workers,
			expires_at = EXCLUDED.expires_at`,
		dateStr, computed.ShareCount, topShareJSON, computed.PoolHashrate, workersJSON, expires)
	if err != nil {
		return nil, fmt.Errorf("failed to insert daily stats: %w", err)
	}

	return computed, nil
}

func ClearDailyStats(db *pgxpool.Pool) error {
	ctx := context.Background()
	if _, err := db.Exec(ctx, `DELETE FROM daily_stats`); err != nil {
		return fmt.Errorf("failed to clear DailyStats: %w", err)
	}
	return nil
}

func DeleteDailyStatsForDay(db *pgxpool.Pool, date time.Time) error {
	ctx := context.Background()
	dateStr := date.Format(time.DateOnly)
	if _, err := db.Exec(ctx, `DELETE FROM daily_stats WHERE date = $1::date`, dateStr); err != nil {
		return fmt.Errorf("failed to delete daily stats for %s: %w", dateStr, err)
	}
	return nil
}

// ShareRowValues returns ordered column values for COPY/batch insert.
func ShareRowValues(s models.ShareLog) ([]any, error) {
	return shareValues(s)
}

// ApplySchema runs the embedded schema against an existing pool.
func ApplySchema(db *pgxpool.Pool) error {
	_, err := db.Exec(context.Background(), schemaSQL)
	return err
}

// CountShares returns the number of rows in shares.
func CountShares(db *pgxpool.Pool) (int64, error) {
	var n int64
	err := db.QueryRow(context.Background(), `SELECT COUNT(*) FROM shares`).Scan(&n)
	return n, err
}