bybit_btcusdt_ingest/utils/stats/stats.go

package stats

import (
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	_ "modernc.org/sqlite"
)

// Colors structure for ANSI output control
type Colors struct {
	Reset  string
	Green  string
	Yellow string
	Cyan   string
	Red    string
	Gray   string
	Bold   string
}

func newColors(enabled bool) Colors {
	if !enabled {
		return Colors{}
	}
	return Colors{
		Reset:  "\033[0m",
		Green:  "\033[32m",
		Yellow: "\033[33m",
		Cyan:   "\033[36m",
		Red:    "\033[31m",
		Gray:   "\033[90m",
		Bold:   "\033[1m",
	}
}

type TableStats struct {
	TableName string
	RowCount  int64
	MinTs     int64
	MaxTs     int64
}

type DBInfo struct {
	Path   string
	Size   int64
	Tables []TableStats
	Error  error
}

// PrintDashboard scans dataDir for multi-stream databases and prints a formatted summary.
func PrintDashboard(dataDir string, noColor bool) {
	c := newColors(!noColor)

	fmt.Printf("%s%s--- Database Dashboard (%s) ---%s\n", c.Bold, c.Green, time.Now().Format(time.RFC822), c.Reset)

	if _, err := os.Stat(dataDir); os.IsNotExist(err) {
		fmt.Printf("  %sData directory not found: %s%s\n", c.Red, dataDir, c.Reset)
		return
	}

	entries, err := os.ReadDir(dataDir)
	if err != nil {
		fmt.Printf("  %sError reading data directory: %v%s\n", c.Red, err, c.Reset)
		return
	}

	var streamDirs []string
	var hasLegacyRootDBs bool

	for _, entry := range entries {
		if entry.IsDir() {
			if entry.Name() != "archive" {
				streamDirs = append(streamDirs, entry.Name())
			}
		} else if strings.HasSuffix(entry.Name(), ".db") {
			hasLegacyRootDBs = true
		}
	}

	sort.Strings(streamDirs)

	var grandTotalSize int64
	var grandTotalRows int64
	var grandTotalArchives int

	// Process legacy root files if present
	if hasLegacyRootDBs {
		size, rows, archives := printStreamDashboard("LEGACY (ROOT)", dataDir, c)
		grandTotalSize += size
		grandTotalRows += rows
		grandTotalArchives += archives
	}

	// Process each stream directory
	for _, stream := range streamDirs {
		streamPath := filepath.Join(dataDir, stream)
		size, rows, archives := printStreamDashboard(strings.ToUpper(stream), streamPath, c)
		grandTotalSize += size
		grandTotalRows += rows
		grandTotalArchives += archives
	}

	// Overall Dashboard Summary
	fmt.Printf("\n%s%s=== GRAND TOTAL DASHBOARD SUMMARY ===%s\n", c.Bold, c.Green, c.Reset)
	fmt.Printf("  %-22s %s%s%s\n", "Total Disk Usage:", c.Yellow, formatBytes(grandTotalSize), c.Reset)
	fmt.Printf("  %-22s %s%d rows%s\n", "Total Records:", c.Yellow, grandTotalRows, c.Reset)
	fmt.Printf("  %-22s %s%d files%s\n", "Total Archive Files:", c.Yellow, grandTotalArchives, c.Reset)
}

func printStreamDashboard(label string, dir string, c Colors) (totalSize int64, totalRows int64, archiveCount int) {
	fmt.Printf("\n%s%s[ STREAM: %s ]%s\n", c.Bold, c.Cyan, label, c.Reset)

	entries, err := os.ReadDir(dir)
	if err != nil {
		fmt.Printf("  %sError reading stream directory %s: %v%s\n", c.Red, dir, err, c.Reset)
		return 0, 0, 0
	}

	var hotDBs []string
	var featureDB string
	var hasArchive bool

	for _, entry := range entries {
		if entry.IsDir() {
			if entry.Name() == "archive" {
				hasArchive = true
			}
			continue
		}

		name := entry.Name()
		if strings.HasSuffix(name, ".db") {
			if name == "features.db" {
				featureDB = filepath.Join(dir, name)
			} else {
				hotDBs = append(hotDBs, filepath.Join(dir, name))
			}
		}
	}

	sort.Strings(hotDBs)

	// 1. Hot DBs
	if len(hotDBs) > 0 {
		fmt.Printf("  %s--- Hot Storage ---%s\n", c.Bold, c.Reset)
		for _, dbPath := range hotDBs {
			size, rows := printDBReport(dbPath, c)
			totalSize += size
			totalRows += rows
		}
	}

	// 2. Feature DB
	if featureDB != "" {
		fmt.Printf("  %s--- Feature Storage ---%s\n", c.Bold, c.Reset)
		size, rows := printDBReport(featureDB, c)
		totalSize += size
		totalRows += rows
	}

	// 3. Archives
	archiveDir := filepath.Join(dir, "archive")
	if hasArchive {
		files, _ := filepath.Glob(filepath.Join(archiveDir, "*.db"))
		archiveCount = len(files)
		if archiveCount > 0 {
			fmt.Printf("  %s--- Archive Storage (%d files) ---%s\n", c.Bold, archiveCount, c.Reset)
			var totalArchivedRows int64
			var totalArchivedSize int64

			for _, f := range files {
				info := inspectDB(f)
				totalArchivedSize += info.Size

				if info.Error != nil {
					fmt.Printf("    %s!%s %-25s: Error reading (%v)\n", c.Red, c.Reset, filepath.Base(f), info.Error)
					continue
				}

				for _, t := range info.Tables {
					totalArchivedRows += t.RowCount
					fmt.Printf("    %s•%s %-25s %s%10s%s | %s%12d rows%s [%s] | %s to %s\n",
						c.Cyan, c.Reset, filepath.Base(f),
						c.Yellow, formatBytes(info.Size), c.Reset,
						c.Gray, t.RowCount, c.Reset, t.TableName,
						formatMs(t.MinTs), formatMs(t.MaxTs))
				}
			}
			fmt.Printf("    %sSubtotal Archived: %d rows across %s%s\n", c.Bold, totalArchivedRows, formatBytes(totalArchivedSize), c.Reset)
			totalSize += totalArchivedSize
			totalRows += totalArchivedRows
		} else {
			fmt.Printf("  %s--- Archive Storage ---%s\n  %sNo archive files found.%s\n", c.Bold, c.Reset, c.Gray, c.Reset)
		}
	}

	fmt.Printf("  %sStream Total: %d rows (%s)%s\n", c.Bold, totalRows, formatBytes(totalSize), c.Reset)
	return totalSize, totalRows, archiveCount
}

func printDBReport(dbPath string, c Colors) (int64, int64) {
	info := inspectDB(dbPath)
	filename := filepath.Base(dbPath)

	fmt.Printf("  %s• Database:%s %s (%s%s%s)\n", c.Cyan, c.Reset, filename, c.Yellow, formatBytes(info.Size), c.Reset)
	if info.Error != nil {
		fmt.Printf("    %sDatabase Error: %v%s\n", c.Red, info.Error, c.Reset)
		return info.Size, 0
	}

	var dbTotalRows int64
	for _, t := range info.Tables {
		dbTotalRows += t.RowCount
		fmt.Printf("    %-18s %s%d%s rows (table: %s%s%s)\n", "Rows:", c.Yellow, t.RowCount, c.Reset, c.Bold, t.TableName, c.Reset)
		if t.RowCount > 0 {
			fmt.Printf("    %-18s %s -> %s\n", "Time Range:", formatMs(t.MinTs), formatMs(t.MaxTs))
			span := formatSpan(t.MinTs, t.MaxTs)
			fmt.Printf("    %-18s %s%s%s\n", "Data Span:", c.Green, span, c.Reset)
		}
	}
	return info.Size, dbTotalRows
}

func inspectDB(dbPath string) DBInfo {
	size := getFullDBSize(dbPath)
	dsn := fmt.Sprintf("file:%s?mode=ro&_journal_mode=WAL", dbPath)
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return DBInfo{Path: dbPath, Size: size, Error: err}
	}
	defer db.Close()

	rows, err := db.Query("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
	if err != nil {
		return DBInfo{Path: dbPath, Size: size, Error: err}
	}
	defer rows.Close()

	var tables []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err == nil {
			tables = append(tables, name)
		}
	}
	rows.Close()

	var tableStats []TableStats
	for _, table := range tables {
		tsCol := findTimestampCol(db, table)
		var count, minTs, maxTs int64
		if tsCol != "" {
			query := fmt.Sprintf("SELECT COUNT(*), COALESCE(MIN(%s), 0), COALESCE(MAX(%s), 0) FROM %s", tsCol, tsCol, table)
			_ = db.QueryRow(query).Scan(&count, &minTs, &maxTs)
		} else {
			query := fmt.Sprintf("SELECT COUNT(*) FROM %s", table)
			_ = db.QueryRow(query).Scan(&count)
		}
		tableStats = append(tableStats, TableStats{
			TableName: table,
			RowCount:  count,
			MinTs:     minTs,
			MaxTs:     maxTs,
		})
	}

	return DBInfo{
		Path:   dbPath,
		Size:   size,
		Tables: tableStats,
	}
}

func findTimestampCol(db *sql.DB, tableName string) string {
	rows, err := db.Query(fmt.Sprintf("PRAGMA table_info(%s)", tableName))
	if err != nil {
		return ""
	}
	defer rows.Close()

	var cols []string
	for rows.Next() {
		var cid int
		var name, typeStr string
		var notNull, pk int
		var dfltValue interface{}
		if err := rows.Scan(&cid, &name, &typeStr, &notNull, &dfltValue, &pk); err == nil {
			cols = append(cols, name)
		}
	}

	priority := []string{"trade_ts", "timestamp", "start_time", "recv_ts", "message_ts"}
	for _, p := range priority {
		for _, c := range cols {
			if strings.EqualFold(c, p) {
				return c
			}
		}
	}

	for _, c := range cols {
		lower := strings.ToLower(c)
		if strings.Contains(lower, "ts") || strings.Contains(lower, "time") {
			return c
		}
	}

	return ""
}

func getFullDBSize(path string) int64 {
	var total int64
	extensions := []string{"", "-wal", "-shm"}
	for _, ext := range extensions {
		if info, err := os.Stat(path + ext); err == nil {
			total += info.Size()
		}
	}
	return total
}

func formatBytes(b int64) string {
	const unit = 1024
	if b < unit {
		return fmt.Sprintf("%d B", b)
	}
	div, exp := int64(unit), 0
	for n := b / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.2f %cB", float64(b)/float64(div), "KMGTPE"[exp])
}

func formatMs(ms int64) string {
	if ms == 0 {
		return "N/A"
	}
	return time.UnixMilli(ms).UTC().Format("2006-01-02 15:04:05 MST")
}

func formatSpan(minTs, maxTs int64) string {
	if minTs == 0 || maxTs == 0 || maxTs <= minTs {
		return "N/A"
	}
	span := time.UnixMilli(maxTs).Sub(time.UnixMilli(minTs)).Round(time.Second)
	return span.String()
}