vctp2/internal/report/snapshots.go

package report

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

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

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

type SnapshotMigrationStats struct {
	HourlyRenamed     int
	HourlyRegistered  int
	DailyRegistered   int
	MonthlyRegistered int
	Errors            int
}

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 MigrateSnapshotRegistry(ctx context.Context, database db.Database) (SnapshotMigrationStats, error) {
	stats := SnapshotMigrationStats{}
	if err := EnsureSnapshotRegistry(ctx, database); err != nil {
		return stats, err
	}

	dbConn := database.DB()
	if _, err := dbConn.ExecContext(ctx, `DELETE FROM snapshot_registry`); err != nil {
		return stats, fmt.Errorf("unable to clear snapshot registry: %w", err)
	}

	allTables, err := ListTablesByPrefix(ctx, database, "inventory_")
	if err != nil {
		return stats, err
	}
	tableSet := make(map[string]struct{}, len(allTables))
	for _, table := range allTables {
		tableSet[table] = struct{}{}
	}

	hourlyTables, err := ListTablesByPrefix(ctx, database, "inventory_hourly_")
	if err != nil {
		return stats, err
	}
	for _, table := range hourlyTables {
		snapshotTime, err := latestSnapshotTime(ctx, dbConn, table)
		if err != nil {
			stats.Errors++
			continue
		}
		if snapshotTime.IsZero() {
			suffix := strings.TrimPrefix(table, "inventory_hourly_")
			if parsed, parseErr := time.Parse("2006010215", suffix); parseErr == nil {
				snapshotTime = parsed
			} else if epoch, parseErr := strconv.ParseInt(suffix, 10, 64); parseErr == nil {
				snapshotTime = time.Unix(epoch, 0)
			}
		}
		if snapshotTime.IsZero() {
			stats.Errors++
			continue
		}

		newName := fmt.Sprintf("inventory_hourly_%d", snapshotTime.Unix())
		if newName != table {
			if _, exists := tableSet[newName]; exists {
				stats.Errors++
				continue
			}
			if err := renameTable(ctx, dbConn, table, newName); err != nil {
				stats.Errors++
				continue
			}
			delete(tableSet, table)
			tableSet[newName] = struct{}{}
			table = newName
			stats.HourlyRenamed++
		}

		if err := RegisterSnapshot(ctx, database, "hourly", table, snapshotTime); err != nil {
			stats.Errors++
			continue
		}
		stats.HourlyRegistered++
	}

	dailyTables, err := ListTablesByPrefix(ctx, database, "inventory_daily_summary_")
	if err != nil {
		return stats, err
	}
	for _, table := range dailyTables {
		suffix := strings.TrimPrefix(table, "inventory_daily_summary_")
		parsed, err := time.Parse("20060102", suffix)
		if err != nil {
			stats.Errors++
			continue
		}
		if err := RegisterSnapshot(ctx, database, "daily", table, parsed); err != nil {
			stats.Errors++
			continue
		}
		stats.DailyRegistered++
	}

	monthlyTables, err := ListTablesByPrefix(ctx, database, "inventory_monthly_summary_")
	if err != nil {
		return stats, err
	}
	for _, table := range monthlyTables {
		suffix := strings.TrimPrefix(table, "inventory_monthly_summary_")
		parsed, err := time.Parse("200601", suffix)
		if err != nil {
			stats.Errors++
			continue
		}
		if err := RegisterSnapshot(ctx, database, "monthly", table, parsed); err != nil {
			stats.Errors++
			continue
		}
		stats.MonthlyRegistered++
	}

	if stats.Errors > 0 {
		return stats, fmt.Errorf("migration completed with %d error(s)", stats.Errors)
	}
	return stats, nil
}

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 ListSnapshotsByRange(ctx context.Context, database db.Database, snapshotType string, start time.Time, end time.Time) ([]SnapshotRecord, error) {
	dbConn := database.DB()
	driver := strings.ToLower(dbConn.DriverName())

	startUnix := start.Unix()
	endUnix := end.Unix()

	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 = ?
  AND snapshot_time >= ?
  AND snapshot_time < ?
ORDER BY snapshot_time ASC, table_name ASC
`, snapshotType, startUnix, endUnix)
	case "pgx", "postgres":
		rows, err = dbConn.QueryxContext(ctx, `
SELECT table_name, snapshot_time, snapshot_type
FROM snapshot_registry
WHERE snapshot_type = $1
  AND snapshot_time >= $2
  AND snapshot_time < $3
ORDER BY snapshot_time ASC, table_name ASC
`, snapshotType, startUnix, endUnix)
	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:04")
	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)
	}

	humanizeTimes := strings.HasPrefix(tableName, "inventory_daily_summary_") || strings.HasPrefix(tableName, "inventory_monthly_summary_")
	type columnSpec struct {
		Name        string
		SourceIndex int
		Humanize    bool
	}
	specs := make([]columnSpec, 0, len(columns)+2)
	for i, columnName := range columns {
		specs = append(specs, columnSpec{Name: columnName, SourceIndex: i})
		if humanizeTimes && columnName == "CreationTime" {
			specs = append(specs, columnSpec{Name: "CreationTimeReadable", SourceIndex: i, Humanize: true})
		}
		if humanizeTimes && columnName == "DeletionTime" {
			specs = append(specs, columnSpec{Name: "DeletionTimeReadable", SourceIndex: i, Humanize: true})
		}
	}

	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, spec := range specs {
		cell := fmt.Sprintf("%s1", string(rune('A'+i)))
		xlsx.SetCellValue(sheetName, cell, spec.Name)
	}

	if endCell, err := excelize.CoordinatesToCellName(len(specs), 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, spec := range specs {
			cell := fmt.Sprintf("%s%d", string(rune('A'+colIndex)), rowIndex)
			value := values[spec.SourceIndex]
			if spec.Humanize {
				xlsx.SetCellValue(sheetName, cell, formatEpochHuman(value))
			} else {
				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
	}
}

func formatEpochHuman(value interface{}) string {
	var epoch int64
	switch v := value.(type) {
	case nil:
		return ""
	case int64:
		epoch = v
	case int32:
		epoch = int64(v)
	case int:
		epoch = int64(v)
	case float64:
		epoch = int64(v)
	case []byte:
		parsed, err := strconv.ParseInt(string(v), 10, 64)
		if err != nil {
			return ""
		}
		epoch = parsed
	case string:
		parsed, err := strconv.ParseInt(v, 10, 64)
		if err != nil {
			return ""
		}
		epoch = parsed
	default:
		return ""
	}
	if epoch <= 0 {
		return ""
	}
	return time.Unix(epoch, 0).Local().Format("Mon 02 Jan 2006 15:04:05 MST")
}

func renameTable(ctx context.Context, dbConn *sqlx.DB, oldName string, newName string) error {
	if err := validateTableName(oldName); err != nil {
		return err
	}
	if err := validateTableName(newName); err != nil {
		return err
	}
	_, err := dbConn.ExecContext(ctx, fmt.Sprintf(`ALTER TABLE %s RENAME TO %s`, oldName, newName))
	if err != nil {
		return fmt.Errorf("failed to rename table %s to %s: %w", oldName, newName, err)
	}
	return nil
}

func latestSnapshotTime(ctx context.Context, dbConn *sqlx.DB, tableName string) (time.Time, error) {
	if err := validateTableName(tableName); err != nil {
		return time.Time{}, err
	}
	query := fmt.Sprintf(`SELECT MAX("SnapshotTime") FROM %s`, tableName)
	var maxTime sql.NullInt64
	if err := dbConn.GetContext(ctx, &maxTime, query); err != nil {
		return time.Time{}, err
	}
	if !maxTime.Valid || maxTime.Int64 <= 0 {
		return time.Time{}, nil
	}
	return time.Unix(maxTime.Int64, 0), nil
}