database.go

package main

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

	_ "github.com/lib/pq"
)

// DBManager handles database operations
type DBManager struct {
	db *sql.DB
}

// NewDBManager creates a new database manager
func NewDBManager(dsn string) (*DBManager, error) {
	// Connect to database - PostgreSQL uses a different connection string format
	db, err := sql.Open("postgres", dsn)
	if err != nil {
		return nil, fmt.Errorf("error connecting to database: %w", err)
	}

	// Test database connection
	if err := db.Ping(); err != nil {
		db.Close()
		return nil, fmt.Errorf("error pinging database: %w", err)
	}

	// Configure connection pool
	db.SetMaxOpenConns(25)
	db.SetMaxIdleConns(5)
	db.SetConnMaxLifetime(5 * time.Minute)

	return &DBManager{
		db: db,
	}, nil
}

// Creates necessary tables if they don't exist
func (m *DBManager) InitTables() error {
	// Create vms table
	_, err := m.db.Exec(`
		CREATE TABLE IF NOT EXISTS vms (
			id SERIAL PRIMARY KEY,
			exoscale_id VARCHAR(36),
			user_id VARCHAR(255) NOT NULL,
			name VARCHAR(255) NOT NULL,
			ip VARCHAR(45),
			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			expires_at TIMESTAMP,
			status VARCHAR(20) DEFAULT 'active'
		)
	`)
	if err != nil {
		return fmt.Errorf("error creating vms table: %w", err)
	}

	// Create user_cooldowns table
	_, err = m.db.Exec(`
	CREATE TABLE IF NOT EXISTS user_cooldowns (
		id SERIAL PRIMARY KEY,
		user_id VARCHAR(255) NOT NULL,
		expires_at TIMESTAMP NOT NULL,
		reason VARCHAR(50) NOT NULL DEFAULT 'session_timeout',
		created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
		UNIQUE (user_id, reason)
	)
	`)
	if err != nil {
		return fmt.Errorf("error creating user_cooldowns table: %w", err)
	}

	log.Println("database tables initialized successfully")
	return nil
}

// Init DB entry when VM creation step has started, initial status is VMStatusCreating
func (m *DBManager) initializeVMCreation(ctx context.Context, userId string) (*VM, error) {
	vm := &VM{
		Name:   fmt.Sprintf("cka-%s", GenerateUniqueID(8)),
		Status: VMStatusCreating,
		UserID: userId,
	}

	_, err := m.db.ExecContext(ctx, `
        INSERT INTO vms (name, user_id, status, created_at)
        VALUES ($1, $2, $3, $4)
    `, vm.Name, vm.UserID, vm.Status, time.Now())

	if err != nil {
		return nil, fmt.Errorf("failed to insert VM record: %w", err)
	}

	return vm, nil
}

// Updates VM information
func (m *DBManager) UpdateVM(ctx context.Context, vm *VM) error {
	_, err := m.db.ExecContext(ctx, `
        UPDATE vms 
        SET status = $1, 
		    exoscale_id = $2,
            ip = $3, 
            expires_at = $4
        WHERE id = $5
    `, vm.Status, vm.ExoscaleID, vm.IP, vm.ExpiresAt, vm.ID)
	return err
}

// Updates the status of a VM in the database
func (m *DBManager) UpdateVMStatus(ctx context.Context, vmID string, status string) error {
	log.Printf("Updating VM %s status to %s", vmID, status)

	result, err := m.db.ExecContext(ctx, `
		UPDATE vms
		SET status = $1
		WHERE id = $2
	`, status, vmID)
	if err != nil {
		return fmt.Errorf("error updating VM status: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("error getting rows affected: %w", err)
	}

	if rowsAffected == 0 {
		return fmt.Errorf("VM with ID %s not found in database", vmID)
	}

	log.Printf("successfully updated VM %s status to %s", vmID, status)
	return nil
}

// Retrieves VMs that have passed their expiration time
func (m *DBManager) GetExpiredVMs(ctx context.Context) ([]*VM, error) {
	rows, err := m.db.QueryContext(ctx, `
		SELECT v.id, v.name, v.ip, v.exoscale_id, v.user_id
		FROM vms v
		WHERE v.expires_at <= $1
		AND v.status = $2
	`, time.Now(), VMStatusActive)
	if err != nil {
		return nil, fmt.Errorf("error querying expired VMs: %w", err)
	}
	defer rows.Close()

	var vms []*VM
	var ipNull sql.NullString
	var exoscaleIDNull sql.NullString

	for rows.Next() {
		var vm VM
		var userID sql.NullString
		if err := rows.Scan(&vm.ID, &vm.Name, &ipNull, &exoscaleIDNull, &userID); err != nil {
			return nil, fmt.Errorf("error scanning VM row: %w", err)
		}

		if userID.Valid {
			vm.UserID = userID.String
		}

		// Set IP from NullString
		if ipNull.Valid {
			vm.IP = ipNull.String
		} else {
			vm.IP = ""
		}

		// Set ExoscaleID from NullString
		if exoscaleIDNull.Valid {
			vm.ExoscaleID = exoscaleIDNull.String
		} else {
			vm.ExoscaleID = ""
		}

		vm.Status = VMStatusActive

		vms = append(vms, &vm)
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("error iterating VM rows: %w", err)
	}

	if len(vms) > 0 {
		log.Printf("found %d expired VMs", len(vms))
	}

	return vms, nil
}

// Retrieves the number of VMs with status VMStatusActive or VMStatusCreating
func (m *DBManager) GetNumberOfActiveVMs(ctx context.Context, vmtype string) (int, error) {
	var count int
	pattern := vmtype + "-%"

	err := m.db.QueryRowContext(ctx, `
		SELECT COUNT(*)
		FROM vms v
		WHERE (v.status = $1 OR v.status = $2)
		AND v.name LIKE $3
	`, VMStatusActive, VMStatusCreating, pattern).Scan(&count)

	if err != nil {
		return 0, fmt.Errorf("error querying active VMs count: %w", err)
	}

	log.Printf("found %d active VMs of type %s", count, vmtype)

	return count, nil
}

// Retrieves the VM associated with a specific user
func (m *DBManager) GetUserVM(ctx context.Context, userID string) (*VM, error) {
	var vm VM
	var status string
	var ipNull sql.NullString
	var expiresAtNull sql.NullTime

	err := m.db.QueryRowContext(ctx, `
        SELECT v.id, v.name, v.user_id, v.ip, v.status, v.expires_at
        FROM vms v
        WHERE v.user_id = $1
        AND (v.status = $2 OR v.status = $3)
        LIMIT 1
    `, userID, VMStatusCreating, VMStatusActive).Scan(&vm.ID, &vm.Name, &vm.UserID, &ipNull, &status, &expiresAtNull)

	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, fmt.Errorf("error querying user VM: %w", err)
	}

	// Set IP from NullString
	if ipNull.Valid {
		vm.IP = ipNull.String
	} else {
		vm.IP = ""
	}

	// Set ExpiresAt from NullTime
	if expiresAtNull.Valid {
		vm.ExpiresAt = expiresAtNull.Time
	} else {
		vm.ExpiresAt = time.Time{}
	}

	vm.UserID = userID
	vm.Status = status

	return &vm, nil
}

// Marks a VM as deleted in the database
func (m *DBManager) MarkVMDeleted(ctx context.Context, vmID string) error {
	return m.UpdateVMStatus(ctx, vmID, VMStatusDeleted)
}

// Checks if a user is in cooldown period
func (m *DBManager) CheckUserCooldown(ctx context.Context, userID string) (bool, map[string]interface{}, error) {
	// Check if the user has any cooldown entry in the database
	var expiryTime time.Time
	var reason string

	err := m.db.QueryRowContext(ctx, `
		SELECT expires_at, reason 
		FROM user_cooldowns 
		WHERE user_id = $1 AND expires_at > $2
		ORDER BY expires_at DESC 
		LIMIT 1
	`, userID, time.Now()).Scan(&expiryTime, &reason)

	if err != nil {
		if err == sql.ErrNoRows {
			// No active cooldown
			return false, nil, nil
		}
		return false, nil, fmt.Errorf("error checking user cooldown: %w", err)
	}

	// User is in cooldown, format the info
	formattedTime := fmt.Sprintf("%s (GMT)", expiryTime.Format("15:04"))

	cooldownInfo := map[string]interface{}{
		"expiryTimestamp": expiryTime.UnixMilli(),
		"formattedTime":   formattedTime,
		"reason":          reason,
	}

	return true, cooldownInfo, nil
}

// Sets a cooldown period for a user
func (m *DBManager) SetUserCooldown(ctx context.Context, userID string, expiresAt time.Time, reason string) error {
	log.Printf("setting cooldown for user %s until %s for reason: %s", userID, expiresAt, reason)

	// Using upsert to handle the case where a cooldown already exists
	_, err := m.db.ExecContext(ctx, `
		INSERT INTO user_cooldowns (user_id, expires_at, reason)
		VALUES ($1, $2, $3)
		ON CONFLICT (user_id, reason) 
		DO UPDATE SET expires_at = $2, created_at = CURRENT_TIMESTAMP
	`, userID, expiresAt, reason)

	if err != nil {
		return fmt.Errorf("error setting user cooldown: %w", err)
	}

	return nil
}

// Deletes expired cooldowns
func (m *DBManager) deleteExpiredCooldowns(ctx context.Context) error {
	log.Printf("deleting expired cooldowns...")

	result, err := m.db.ExecContext(ctx, `
        DELETE FROM user_cooldowns 
        WHERE expires_at < $1
    `, time.Now())

	if err != nil {
		return fmt.Errorf("error deleting expired cooldowns: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("error getting rows affected for expired cooldowns: %w", err)
	}

	if rowsAffected > 0 {
		log.Printf("deleted %d expired user cooldowns", rowsAffected)
	}

	return nil
}

// Closes the database connection
func (m *DBManager) Close() error {
	if m.db != nil {
		return m.db.Close()
	}
	return nil
}