whatshooked/pkg/storage/db.go

package storage

import (
	"context"
	"database/sql"
	"fmt"
	"time"

	"git.warky.dev/wdevs/whatshooked/pkg/config"
	"git.warky.dev/wdevs/whatshooked/pkg/models"
	"github.com/uptrace/bun"
	"github.com/uptrace/bun/dialect/pgdialect"
	"github.com/uptrace/bun/dialect/sqlitedialect"
	"github.com/uptrace/bun/driver/pgdriver"
	"github.com/uptrace/bun/driver/sqliteshim"
	"github.com/uptrace/bun/extra/bundebug"
)

// DB is the global database instance
var DB *bun.DB
var dbType string // Store the database type for later use

// Initialize sets up the database connection based on configuration
func Initialize(cfg *config.DatabaseConfig) error {
	var sqldb *sql.DB
	var err error

	dbType = cfg.Type
	switch cfg.Type {
	case "postgres", "postgresql":
		dsn := fmt.Sprintf("postgres://%s:%s@%s:%d/%s?sslmode=disable",
			cfg.Username, cfg.Password, cfg.Host, cfg.Port, cfg.Database)
		sqldb = sql.OpenDB(pgdriver.NewConnector(pgdriver.WithDSN(dsn)))
		DB = bun.NewDB(sqldb, pgdialect.New())

	case "sqlite":
		sqldb, err = sql.Open(sqliteshim.ShimName, cfg.SQLitePath)
		if err != nil {
			return fmt.Errorf("failed to open sqlite database: %w", err)
		}
		DB = bun.NewDB(sqldb, sqlitedialect.New())

	default:
		return fmt.Errorf("unsupported database type: %s", cfg.Type)
	}

	// Add query hook for debugging (optional, can be removed in production)
	DB.AddQueryHook(bundebug.NewQueryHook(
		bundebug.WithVerbose(true),
		bundebug.FromEnv("BUNDEBUG"),
	))

	// Set connection pool settings
	sqldb.SetMaxIdleConns(10)
	sqldb.SetMaxOpenConns(100)
	sqldb.SetConnMaxLifetime(time.Hour)

	// Test the connection
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := sqldb.PingContext(ctx); err != nil {
		return fmt.Errorf("failed to ping database: %w", err)
	}

	return nil
}

// CreateTables creates database tables based on BUN models
func CreateTables(ctx context.Context) error {
	if DB == nil {
		return fmt.Errorf("database not initialized")
	}

	// For SQLite, use raw SQL with compatible syntax
	if dbType == "sqlite" {
		return createTablesSQLite(ctx)
	}

	// For PostgreSQL, use BUN's auto-generation
	models := []interface{}{
		(*models.ModelPublicUsers)(nil),
		(*models.ModelPublicAPIKey)(nil),
		(*models.ModelPublicHook)(nil),
		(*models.ModelPublicWhatsappAccount)(nil),
		(*models.ModelPublicEventLog)(nil),
		(*models.ModelPublicSession)(nil),
	}

	for _, model := range models {
		_, err := DB.NewCreateTable().Model(model).IfNotExists().Exec(ctx)
		if err != nil {
			return fmt.Errorf("failed to create table: %w", err)
		}
	}

	if err := ensureMessageCacheTable(ctx); err != nil {
		return err
	}

	return nil
}

// createTablesSQLite creates tables using SQLite-compatible SQL
func createTablesSQLite(ctx context.Context) error {
	tables := []string{
		// Users table
		`CREATE TABLE IF NOT EXISTS users (
			id VARCHAR(36) PRIMARY KEY,
			username VARCHAR(255) NOT NULL UNIQUE,
			email VARCHAR(255) NOT NULL UNIQUE,
			password VARCHAR(255) NOT NULL,
			full_name VARCHAR(255),
			role VARCHAR(50) NOT NULL DEFAULT 'user',
			active BOOLEAN NOT NULL DEFAULT 1,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			deleted_at TIMESTAMP
		)`,

		// API Keys table
		`CREATE TABLE IF NOT EXISTS api_keys (
			id VARCHAR(36) PRIMARY KEY,
			user_id VARCHAR(36) NOT NULL,
			name VARCHAR(255) NOT NULL,
			key VARCHAR(255) NOT NULL UNIQUE,
			key_prefix VARCHAR(20),
			permissions TEXT,
			active BOOLEAN NOT NULL DEFAULT 1,
			last_used_at TIMESTAMP,
			expires_at TIMESTAMP,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			deleted_at TIMESTAMP,
			FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
		)`,

		// Hooks table
		`CREATE TABLE IF NOT EXISTS hooks (
			id VARCHAR(36) PRIMARY KEY,
			user_id VARCHAR(36) NOT NULL,
			name VARCHAR(255) NOT NULL,
			url TEXT NOT NULL,
			method VARCHAR(10) NOT NULL DEFAULT 'POST',
			headers TEXT,
			events TEXT,
			active BOOLEAN NOT NULL DEFAULT 1,
			allow_insecure BOOLEAN NOT NULL DEFAULT 0,
			retry_count INTEGER NOT NULL DEFAULT 3,
			timeout INTEGER NOT NULL DEFAULT 30,
			secret VARCHAR(255),
			description TEXT,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			deleted_at TIMESTAMP,
			FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
		)`,

		// WhatsApp Accounts table
		`CREATE TABLE IF NOT EXISTS whatsapp_account (
			id VARCHAR(36) PRIMARY KEY,
			user_id VARCHAR(36) NOT NULL,
			phone_number VARCHAR(20) NOT NULL UNIQUE,
			display_name VARCHAR(255),
			account_type VARCHAR(50) NOT NULL DEFAULT 'whatsmeow',
			config TEXT,
			active BOOLEAN NOT NULL DEFAULT 1,
			status VARCHAR(50) NOT NULL DEFAULT 'disconnected',
			session_path TEXT,
			last_connected_at TIMESTAMP,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			deleted_at TIMESTAMP,
			FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
		)`,

		// Event Logs table
		`CREATE TABLE IF NOT EXISTS event_logs (
			id VARCHAR(36) PRIMARY KEY,
			user_id VARCHAR(36),
			account_id VARCHAR(36),
			event_type VARCHAR(100) NOT NULL,
			event_data TEXT,
			from_number VARCHAR(20),
			to_number VARCHAR(20),
			message_id VARCHAR(255),
			status VARCHAR(50),
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE SET NULL,
			FOREIGN KEY (account_id) REFERENCES whatsapp_account(id) ON DELETE SET NULL
		)`,

		// Sessions table
		`CREATE TABLE IF NOT EXISTS sessions (
			id VARCHAR(36) PRIMARY KEY,
			user_id VARCHAR(36) NOT NULL,
			token VARCHAR(500) NOT NULL UNIQUE,
			expires_at TIMESTAMP NOT NULL,
			ip_address VARCHAR(45),
			user_agent TEXT,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
		)`,

		// Message Cache table
		`CREATE TABLE IF NOT EXISTS message_cache (
			id VARCHAR(128) PRIMARY KEY,
			account_id VARCHAR(64) NOT NULL DEFAULT '',
			event_type VARCHAR(100) NOT NULL,
			event_data TEXT NOT NULL,
			message_id VARCHAR(255) NOT NULL DEFAULT '',
			from_number VARCHAR(64) NOT NULL DEFAULT '',
			to_number VARCHAR(64) NOT NULL DEFAULT '',
			reason TEXT NOT NULL DEFAULT '',
			attempts INTEGER NOT NULL DEFAULT 0,
			timestamp TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			last_attempt TIMESTAMP,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
		)`,

		`CREATE INDEX IF NOT EXISTS idx_message_cache_timestamp ON message_cache(timestamp DESC)`,
		`CREATE INDEX IF NOT EXISTS idx_message_cache_event_type ON message_cache(event_type)`,
	}

	for _, sql := range tables {
		if _, err := DB.ExecContext(ctx, sql); err != nil {
			return fmt.Errorf("failed to create table: %w", err)
		}
	}

	if err := migrateLegacyWhatsAppAccountsTable(ctx); err != nil {
		return err
	}
	if err := ensureWhatsAppAccountColumnsSQLite(ctx); err != nil {
		return err
	}

	if err := ensureMessageCacheTable(ctx); err != nil {
		return err
	}

	return nil
}

func migrateLegacyWhatsAppAccountsTable(ctx context.Context) error {
	if DB == nil {
		return fmt.Errorf("database not initialized")
	}

	var legacyCount int
	if err := DB.NewSelect().
		Table("sqlite_master").
		ColumnExpr("COUNT(1)").
		Where("type = 'table' AND name = 'whatsapp_accounts'").
		Scan(ctx, &legacyCount); err != nil {
		return fmt.Errorf("failed to inspect legacy whatsapp_accounts table: %w", err)
	}

	var currentCount int
	if err := DB.NewSelect().
		Table("sqlite_master").
		ColumnExpr("COUNT(1)").
		Where("type = 'table' AND name = 'whatsapp_account'").
		Scan(ctx, &currentCount); err != nil {
		return fmt.Errorf("failed to inspect whatsapp_account table: %w", err)
	}

	if legacyCount > 0 {
		if currentCount == 0 {
			if _, err := DB.ExecContext(ctx, `ALTER TABLE whatsapp_accounts RENAME TO whatsapp_account`); err != nil {
				return fmt.Errorf("failed to migrate table whatsapp_accounts -> whatsapp_account: %w", err)
			}
			return nil
		}

		mergeSQL := `INSERT OR IGNORE INTO whatsapp_account
			(id, user_id, phone_number, display_name, account_type, config, active, status, session_path, last_connected_at, created_at, updated_at, deleted_at)
			SELECT
				id,
				user_id,
				phone_number,
				'' AS display_name,
				COALESCE(account_type, 'whatsmeow') AS account_type,
				COALESCE(business_api_config, '') AS config,
				COALESCE(active, 1) AS active,
				CASE
					WHEN COALESCE(active, 1) = 0 THEN 'disconnected'
					WHEN COALESCE(connected, 0) = 1 THEN 'connected'
					ELSE 'disconnected'
				END AS status,
				'' AS session_path,
				last_connected_at,
				created_at,
				updated_at,
				deleted_at
			FROM whatsapp_accounts`
		if _, err := DB.ExecContext(ctx, mergeSQL); err != nil {
			return fmt.Errorf("failed to merge legacy whatsapp_accounts data: %w", err)
		}
		if _, err := DB.ExecContext(ctx, `DROP TABLE whatsapp_accounts`); err != nil {
			return fmt.Errorf("failed to drop legacy whatsapp_accounts table: %w", err)
		}
	}

	return nil
}

func ensureWhatsAppAccountColumnsSQLite(ctx context.Context) error {
	if err := ensureSQLiteColumn(ctx, "whatsapp_account", "display_name", "VARCHAR(255)"); err != nil {
		return err
	}
	if err := ensureSQLiteColumn(ctx, "whatsapp_account", "config", "TEXT"); err != nil {
		return err
	}
	if err := ensureSQLiteColumn(ctx, "whatsapp_account", "status", "VARCHAR(50) NOT NULL DEFAULT 'disconnected'"); err != nil {
		return err
	}
	if err := ensureSQLiteColumn(ctx, "whatsapp_account", "session_path", "TEXT"); err != nil {
		return err
	}
	if err := ensureSQLiteColumn(ctx, "whatsapp_account", "updated_at", "TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP"); err != nil {
		return err
	}

	// Backfill config/status from legacy columns if they still exist on this table.
	if hasBusinessAPIConfig, err := sqliteColumnExists(ctx, "whatsapp_account", "business_api_config"); err != nil {
		return err
	} else if hasBusinessAPIConfig {
		if _, err := DB.ExecContext(ctx, `UPDATE whatsapp_account SET config = COALESCE(config, business_api_config, '')`); err != nil {
			return fmt.Errorf("failed to backfill config from business_api_config: %w", err)
		}
	}

	if hasConnected, err := sqliteColumnExists(ctx, "whatsapp_account", "connected"); err != nil {
		return err
	} else if hasConnected {
		if _, err := DB.ExecContext(ctx, `UPDATE whatsapp_account
			SET status = CASE
				WHEN COALESCE(active, 1) = 0 THEN 'disconnected'
				WHEN COALESCE(connected, 0) = 1 THEN 'connected'
				WHEN status = '' OR status IS NULL THEN 'disconnected'
				ELSE status
			END`); err != nil {
			return fmt.Errorf("failed to backfill status from connected column: %w", err)
		}
	}

	return nil
}

func sqliteColumnExists(ctx context.Context, table, column string) (bool, error) {
	rows, err := DB.QueryContext(ctx, fmt.Sprintf("PRAGMA table_info(%s)", table))
	if err != nil {
		return false, fmt.Errorf("failed to inspect sqlite table %s: %w", table, err)
	}
	defer rows.Close()

	for rows.Next() {
		var cid int
		var colName string
		var colType string
		var notNull int
		var defaultValue any
		var pk int
		if err := rows.Scan(&cid, &colName, &colType, &notNull, &defaultValue, &pk); err != nil {
			return false, fmt.Errorf("failed to scan sqlite table info for %s: %w", table, err)
		}
		if colName == column {
			return true, nil
		}
	}
	if err := rows.Err(); err != nil {
		return false, fmt.Errorf("failed reading sqlite table info for %s: %w", table, err)
	}
	return false, nil
}

func ensureSQLiteColumn(ctx context.Context, table, name, definition string) error {
	exists, err := sqliteColumnExists(ctx, table, name)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}

	query := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", table, name, definition)
	if _, err := DB.ExecContext(ctx, query); err != nil {
		return fmt.Errorf("failed to add sqlite column %s.%s: %w", table, name, err)
	}
	return nil
}

func ensureMessageCacheTable(ctx context.Context) error {
	if DB == nil {
		return fmt.Errorf("database not initialized")
	}

	if dbType == "postgres" || dbType == "postgresql" {
		queries := []string{
			`CREATE TABLE IF NOT EXISTS message_cache (
				id VARCHAR(128) PRIMARY KEY,
				account_id VARCHAR(64) NOT NULL DEFAULT '',
				event_type VARCHAR(100) NOT NULL,
				event_data JSONB NOT NULL,
				message_id VARCHAR(255) NOT NULL DEFAULT '',
				from_number VARCHAR(64) NOT NULL DEFAULT '',
				to_number VARCHAR(64) NOT NULL DEFAULT '',
				reason TEXT NOT NULL DEFAULT '',
				attempts INTEGER NOT NULL DEFAULT 0,
				timestamp TIMESTAMPTZ NOT NULL,
				last_attempt TIMESTAMPTZ,
				created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
			)`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS account_id VARCHAR(64) NOT NULL DEFAULT ''`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS event_type VARCHAR(100) NOT NULL DEFAULT ''`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS event_data JSONB NOT NULL DEFAULT '{}'::jsonb`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS message_id VARCHAR(255) NOT NULL DEFAULT ''`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS from_number VARCHAR(64) NOT NULL DEFAULT ''`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS to_number VARCHAR(64) NOT NULL DEFAULT ''`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS reason TEXT NOT NULL DEFAULT ''`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS attempts INTEGER NOT NULL DEFAULT 0`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW()`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS last_attempt TIMESTAMPTZ`,
			`ALTER TABLE message_cache ADD COLUMN IF NOT EXISTS created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()`,
			`CREATE INDEX IF NOT EXISTS idx_message_cache_timestamp ON message_cache (timestamp DESC)`,
			`CREATE INDEX IF NOT EXISTS idx_message_cache_event_type ON message_cache (event_type)`,
		}
		for _, query := range queries {
			if _, err := DB.ExecContext(ctx, query); err != nil {
				return fmt.Errorf("failed to ensure postgres message_cache table: %w", err)
			}
		}
		return nil
	}

	queries := []string{
		`CREATE TABLE IF NOT EXISTS message_cache (
			id TEXT PRIMARY KEY,
			account_id TEXT NOT NULL DEFAULT '',
			event_type TEXT NOT NULL,
			event_data TEXT NOT NULL,
			message_id TEXT NOT NULL DEFAULT '',
			from_number TEXT NOT NULL DEFAULT '',
			to_number TEXT NOT NULL DEFAULT '',
			reason TEXT NOT NULL DEFAULT '',
			attempts INTEGER NOT NULL DEFAULT 0,
			timestamp DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
			last_attempt DATETIME,
			created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE INDEX IF NOT EXISTS idx_message_cache_timestamp ON message_cache(timestamp DESC)`,
		`CREATE INDEX IF NOT EXISTS idx_message_cache_event_type ON message_cache(event_type)`,
	}
	for _, query := range queries {
		if _, err := DB.ExecContext(ctx, query); err != nil {
			return fmt.Errorf("failed to ensure sqlite message_cache table: %w", err)
		}
	}

	return nil
}

// Close closes the database connection
func Close() error {
	if DB == nil {
		return nil
	}

	return DB.Close()
}

// GetDB returns the database instance
func GetDB() *bun.DB {
	return DB
}

// HealthCheck checks if the database connection is healthy
func HealthCheck() error {
	if DB == nil {
		return fmt.Errorf("database not initialized")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	return DB.PingContext(ctx)
}