vctp2/internal/report/snapshots.go

package report

import (
	"bytes"
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"strings"
	"time"
	"vctp/db"

	"github.com/jmoiron/sqlx"
	"github.com/xuri/excelize/v2"
)

type SnapshotRecord struct {
	TableName    string
	SnapshotTime time.Time
	SnapshotType string
}

func ListTablesByPrefix(ctx context.Context, database db.Database, prefix string) ([]string, error) {
	dbConn := database.DB()
	driver := strings.ToLower(dbConn.DriverName())
	pattern := prefix + "%"

	var rows *sqlx.Rows
	var err error

	switch driver {
	case "sqlite":
		rows, err = dbConn.QueryxContext(ctx, `
SELECT name
FROM sqlite_master
WHERE type = 'table'
  AND name LIKE ?
ORDER BY name DESC
`, pattern)
	case "pgx", "postgres":
		rows, err = dbConn.QueryxContext(ctx, `
SELECT tablename
FROM pg_catalog.pg_tables
WHERE schemaname = 'public'
  AND tablename LIKE $1
ORDER BY tablename DESC
`, pattern)
	default:
		return nil, fmt.Errorf("unsupported driver for listing tables: %s", driver)
	}

	if err != nil {
		return nil, err
	}
	defer rows.Close()

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

func EnsureSnapshotRegistry(ctx context.Context, database db.Database) error {
	dbConn := database.DB()
	driver := strings.ToLower(dbConn.DriverName())
	switch driver {
	case "sqlite":
		_, err := dbConn.ExecContext(ctx, `
CREATE TABLE IF NOT EXISTS snapshot_registry (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  snapshot_type TEXT NOT NULL,
  table_name TEXT NOT NULL UNIQUE,
  snapshot_time BIGINT NOT NULL
)
`)
		return err
	case "pgx", "postgres":
		_, err := dbConn.ExecContext(ctx, `
CREATE TABLE IF NOT EXISTS snapshot_registry (
  id BIGSERIAL PRIMARY KEY,
  snapshot_type TEXT NOT NULL,
  table_name TEXT NOT NULL UNIQUE,
  snapshot_time BIGINT NOT NULL
)
`)
		return err
	default:
		return fmt.Errorf("unsupported driver for snapshot registry: %s", driver)
	}
}

func RegisterSnapshot(ctx context.Context, database db.Database, snapshotType string, tableName string, snapshotTime time.Time) error {
	if snapshotType == "" || tableName == "" {
		return fmt.Errorf("snapshot type or table name is empty")
	}
	dbConn := database.DB()
	driver := strings.ToLower(dbConn.DriverName())
	switch driver {
	case "sqlite":
		_, err := dbConn.ExecContext(ctx, `
INSERT OR IGNORE INTO snapshot_registry (snapshot_type, table_name, snapshot_time)
VALUES (?, ?, ?)
`, snapshotType, tableName, snapshotTime.Unix())
		return err
	case "pgx", "postgres":
		_, err := dbConn.ExecContext(ctx, `
INSERT INTO snapshot_registry (snapshot_type, table_name, snapshot_time)
VALUES ($1, $2, $3)
ON CONFLICT (table_name) DO NOTHING
`, snapshotType, tableName, snapshotTime.Unix())
		return err
	default:
		return fmt.Errorf("unsupported driver for snapshot registry: %s", driver)
	}
}

func DeleteSnapshotRecord(ctx context.Context, database db.Database, tableName string) error {
	if tableName == "" {
		return nil
	}
	dbConn := database.DB()
	driver := strings.ToLower(dbConn.DriverName())
	switch driver {
	case "sqlite":
		_, err := dbConn.ExecContext(ctx, `DELETE FROM snapshot_registry WHERE table_name = ?`, tableName)
		return err
	case "pgx", "postgres":
		_, err := dbConn.ExecContext(ctx, `DELETE FROM snapshot_registry WHERE table_name = $1`, tableName)
		return err
	default:
		return fmt.Errorf("unsupported driver for snapshot registry: %s", driver)
	}
}

func ListSnapshots(ctx context.Context, database db.Database, snapshotType string) ([]SnapshotRecord, error) {
	dbConn := database.DB()
	driver := strings.ToLower(dbConn.DriverName())

	var rows *sqlx.Rows
	var err error

	switch driver {
	case "sqlite":
		rows, err = dbConn.QueryxContext(ctx, `
SELECT table_name, snapshot_time, snapshot_type
FROM snapshot_registry
WHERE snapshot_type = ?
ORDER BY snapshot_time DESC, table_name DESC
`, snapshotType)
	case "pgx", "postgres":
		rows, err = dbConn.QueryxContext(ctx, `
SELECT table_name, snapshot_time, snapshot_type
FROM snapshot_registry
WHERE snapshot_type = $1
ORDER BY snapshot_time DESC, table_name DESC
`, snapshotType)
	default:
		return nil, fmt.Errorf("unsupported driver for listing snapshots: %s", driver)
	}

	if err != nil {
		return nil, err
	}
	defer rows.Close()

	records := make([]SnapshotRecord, 0)
	for rows.Next() {
		var (
			tableName    string
			snapshotTime int64
			recordType   string
		)
		if err := rows.Scan(&tableName, &snapshotTime, &recordType); err != nil {
			return nil, err
		}
		records = append(records, SnapshotRecord{
			TableName:    tableName,
			SnapshotTime: time.Unix(snapshotTime, 0),
			SnapshotType: recordType,
		})
	}
	return records, rows.Err()
}

func FormatSnapshotLabel(snapshotType string, snapshotTime time.Time, tableName string) string {
	switch snapshotType {
	case "hourly":
		return snapshotTime.Format("2006-01-02 15:00")
	case "daily":
		return snapshotTime.Format("2006-01-02")
	case "monthly":
		return snapshotTime.Format("2006-01")
	default:
		return tableName
	}
}

func CreateTableReport(logger *slog.Logger, Database db.Database, ctx context.Context, tableName string) ([]byte, error) {
	if err := validateTableName(tableName); err != nil {
		return nil, err
	}

	dbConn := Database.DB()
	columns, err := tableColumns(ctx, dbConn, tableName)
	if err != nil {
		return nil, err
	}
	if len(columns) == 0 {
		return nil, fmt.Errorf("no columns found for table %s", tableName)
	}

	query := fmt.Sprintf(`SELECT * FROM %s`, tableName)
	orderBy := snapshotOrderBy(columns)
	if orderBy != "" {
		query = fmt.Sprintf(`%s ORDER BY "%s" DESC`, query, orderBy)
	}

	rows, err := dbConn.QueryxContext(ctx, query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	sheetName := "Snapshot Report"
	var buffer bytes.Buffer

	xlsx := excelize.NewFile()
	if err := xlsx.SetSheetName("Sheet1", sheetName); err != nil {
		return nil, err
	}
	if err := xlsx.SetDocProps(&excelize.DocProperties{
		Creator: "vctp",
		Created: time.Now().Format(time.RFC3339),
	}); err != nil {
		logger.Error("Error setting document properties", "error", err, "sheet_name", sheetName)
	}

	for i, columnName := range columns {
		cell := fmt.Sprintf("%s1", string(rune('A'+i)))
		xlsx.SetCellValue(sheetName, cell, columnName)
	}

	if endCell, err := excelize.CoordinatesToCellName(len(columns), 1); err == nil {
		filterRange := "A1:" + endCell
		if err := xlsx.AutoFilter(sheetName, filterRange, nil); err != nil {
			logger.Error("Error setting autofilter", "error", err)
		}
	}

	headerStyle, err := xlsx.NewStyle(&excelize.Style{
		Font: &excelize.Font{
			Bold: true,
		},
	})
	if err == nil {
		if err := xlsx.SetRowStyle(sheetName, 1, 1, headerStyle); err != nil {
			logger.Error("Error setting header style", "error", err)
		}
	}

	rowIndex := 2
	for rows.Next() {
		values, err := scanRowValues(rows, len(columns))
		if err != nil {
			return nil, err
		}
		for colIndex, value := range values {
			cell := fmt.Sprintf("%s%d", string(rune('A'+colIndex)), rowIndex)
			xlsx.SetCellValue(sheetName, cell, normalizeCellValue(value))
		}
		rowIndex++
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	if err := xlsx.SetPanes(sheetName, &excelize.Panes{
		Freeze:      true,
		Split:       false,
		XSplit:      0,
		YSplit:      1,
		TopLeftCell: "A2",
		ActivePane:  "bottomLeft",
		Selection: []excelize.Selection{
			{SQRef: "A2", ActiveCell: "A2", Pane: "bottomLeft"},
		},
	}); err != nil {
		logger.Error("Error freezing top row", "error", err)
	}

	if err := xlsx.Write(&buffer); err != nil {
		return nil, err
	}
	return buffer.Bytes(), nil
}

func validateTableName(name string) error {
	if name == "" {
		return fmt.Errorf("table name is empty")
	}
	for _, r := range name {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '_' {
			continue
		}
		return fmt.Errorf("invalid table name: %s", name)
	}
	return nil
}

func tableColumns(ctx context.Context, dbConn *sqlx.DB, tableName string) ([]string, error) {
	driver := strings.ToLower(dbConn.DriverName())
	switch driver {
	case "sqlite":
		query := fmt.Sprintf(`PRAGMA table_info("%s")`, tableName)
		rows, err := dbConn.QueryxContext(ctx, query)
		if err != nil {
			return nil, err
		}
		defer rows.Close()

		columns := make([]string, 0)
		for rows.Next() {
			var (
				cid        int
				name       string
				colType    string
				notNull    int
				defaultVal sql.NullString
				pk         int
			)
			if err := rows.Scan(&cid, &name, &colType, &notNull, &defaultVal, &pk); err != nil {
				return nil, err
			}
			columns = append(columns, name)
		}
		return columns, rows.Err()
	case "pgx", "postgres":
		rows, err := dbConn.QueryxContext(ctx, `
SELECT column_name
FROM information_schema.columns
WHERE table_schema = 'public'
  AND table_name = $1
ORDER BY ordinal_position
`, tableName)
		if err != nil {
			return nil, err
		}
		defer rows.Close()

		columns := make([]string, 0)
		for rows.Next() {
			var name string
			if err := rows.Scan(&name); err != nil {
				return nil, err
			}
			columns = append(columns, name)
		}
		return columns, rows.Err()
	default:
		return nil, fmt.Errorf("unsupported driver for table columns: %s", driver)
	}
}

func snapshotOrderBy(columns []string) string {
	normalized := make(map[string]struct{}, len(columns))
	for _, col := range columns {
		normalized[strings.ToLower(col)] = struct{}{}
	}
	if _, ok := normalized["snapshottime"]; ok {
		return "SnapshotTime"
	}
	if _, ok := normalized["samplespresent"]; ok {
		return "SamplesPresent"
	}
	if _, ok := normalized["avgispresent"]; ok {
		return "AvgIsPresent"
	}
	if _, ok := normalized["name"]; ok {
		return "Name"
	}
	return ""
}

func scanRowValues(rows *sqlx.Rows, columnCount int) ([]interface{}, error) {
	rawValues := make([]interface{}, columnCount)
	scanArgs := make([]interface{}, columnCount)
	for i := range rawValues {
		scanArgs[i] = &rawValues[i]
	}
	if err := rows.Scan(scanArgs...); err != nil {
		return nil, err
	}
	return rawValues, nil
}

func normalizeCellValue(value interface{}) interface{} {
	switch v := value.(type) {
	case nil:
		return ""
	case []byte:
		return string(v)
	case time.Time:
		return v.Format(time.RFC3339)
	default:
		return v
	}
}