manager.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"time"

	egoscale "github.com/exoscale/egoscale/v3"
	"github.com/exoscale/egoscale/v3/credentials"
	"github.com/nats-io/nats.go"
)

// VMManager handles VM operations
type VMManager struct {
	config     *Config
	exoClient  *egoscale.Client
	natsClient *nats.Conn
	db         *DBManager
}

// NewVMManager creates a new VM manager
func NewVMManager(cfg *Config) (*VMManager, error) {
	// Create Exoscale client
	creds := credentials.NewStaticCredentials(cfg.ExoscaleAPIKey, cfg.ExoscaleAPISecret)
	client, err := egoscale.NewClient(
		creds,
		egoscale.ClientOptWithEndpoint(getExoscaleEndpoint(cfg.ExoscaleEndpoint)),
	)
	if err != nil {
		return nil, fmt.Errorf("error creating Exoscale client: %w", err)
	}

	// Connect to NATS
	nc, err := nats.Connect(cfg.NatsURL)
	if err != nil {
		return nil, fmt.Errorf("error connecting to NATS: %w", err)
	}

	// Create database manager
	dbManager, err := NewDBManager(cfg.DatabaseDSN)
	if err != nil {
		nc.Close() // Close NATS connection if DB connection fails
		return nil, fmt.Errorf("error creating database manager: %w", err)
	}

	return &VMManager{
		config:     cfg,
		exoClient:  client,
		natsClient: nc,
		db:         dbManager,
	}, nil
}

// InitDatabase initializes the database tables
func (m *VMManager) InitDatabase() error {
	return m.db.InitTables()
}

// Close closes all connections
func (m *VMManager) Close() {
	if m.natsClient != nil {
		m.natsClient.Close()
	}
	if m.db != nil {
		m.db.Close()
	}
	log.Println("VM manager connections closed")
}

// Helper function to map endpoint string to Exoscale endpoint
func getExoscaleEndpoint(endpoint string) egoscale.Endpoint {
	switch endpoint {
	case "CHDk2":
		return egoscale.CHDk2
	case "CHGva2":
		return egoscale.CHGva2
	case "ATVie1":
		return egoscale.ATVie1
	case "ATVie2":
		return egoscale.ATVie2
	case "DEFra1":
		return egoscale.DEFra1
	default:
		// Default to Switzerland
		return egoscale.CHDk2
	}
}

// Sets up NATS subscriptions for VM management
func (m *VMManager) SetupNATSSubscriptions() error {
	// Subscribe to VM creation requests
	_, err := m.natsClient.Subscribe("vm.create", func(msg *nats.Msg) {
		log.Println("received request on vm.create")

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

		// Forget if no Reply subject if provided
		if msg.Reply == "" {
			log.Println("no reply subject provided")
			return
		}

		// Parse the request to get userId
		var request struct {
			UserID      string `json:"userId"`
			RequestID   string `json:"requestId"`
			RequestedAt int64  `json:"requestedAt"`
			VMType      string `json:"vmtype"`
		}

		if err := json.Unmarshal(msg.Data, &request); err != nil {
			log.Printf("failed to unmarshal VM creation request: %v", err)

			errResponse := map[string]string{
				"status": "error",
				"error":  "invalid request format",
			}
			errData, _ := json.Marshal(errResponse)
			m.natsClient.Publish(msg.Reply, errData)
			return
		}

		log.Printf("received [%s] VM creation request for user [%s]", request.VMType, request.UserID)

		// Check if a VM is already associated to this user
		vm, err := m.db.GetUserVM(ctx, request.UserID)
		if err != nil {
			log.Printf("error retrieving VM for user %s: %v", request.UserID, err)
			errResp := map[string]interface{}{
				"status": "error",
				"error":  err.Error(),
			}
			errData, _ := json.Marshal(errResp)
			m.natsClient.Publish(msg.Reply, errData)
			return
		}

		// Send VM information is one is found
		if vm != nil {
			log.Printf("[%s] VM found for user %s", vm.Status, request.UserID)

			var response map[string]interface{}
			response = make(map[string]interface{})

			// Handle VMStatusCreating status
			if vm.Status == VMStatusCreating {
				response["userId"] = request.UserID
				response["vmId"] = vm.ID
				response["message"] = "VM under creation found"
				response["status"] = vm.Status
			}

			// Handle VMStatusActive status
			if vm.Status == VMStatusActive {
				response["userId"] = request.UserID
				response["vmId"] = vm.ID
				response["vmIp"] = vm.IP
				response["message"] = "active VM found"
				response["status"] = vm.Status
			}

			respData, _ := json.Marshal(response)
			m.natsClient.Publish(msg.Reply, respData)
			return
		}

		// No VM found for the user

		// Get configuration of the type of VM requested
		vmConfig, err := getVMConfig(request.VMType)
		if err != nil {
			log.Printf("failed to get config for VM of type %s: %v", request.VMType, err)
		}

		log.Printf("configuration for VM of type %s: %v", request.VMType, vmConfig)

		// Make sure quota is not reached for this type of VM
		activeVMsNumber, err := m.db.GetNumberOfActiveVMs(ctx, request.VMType)
		if err != nil {
			log.Printf("error retrieving number of active VMs of type %s: %v", err, request.VMType)
			errResp := map[string]interface{}{
				"status": "error",
				"error":  err.Error(),
			}
			errData, _ := json.Marshal(errResp)
			m.natsClient.Publish(msg.Reply, errData)
			return
		} else {
			if activeVMsNumber >= vmConfig.Quota {
				log.Printf("Number of active VM (%d) of type (%s) has reached the allowed quota (%d)", activeVMsNumber, request.VMType, m.config.VMMaxNumber)
				errResp := map[string]interface{}{
					"status": "error",
					"error":  "No more environment available. Please try again later.",
				}
				errData, _ := json.Marshal(errResp)
				m.natsClient.Publish(msg.Reply, errData)
				return
			}
		}

		// Create a VM in DB only to keep track of the creation step
		vm, err = m.db.initializeVMCreation(context.Background(), request.UserID)
		if err != nil {
			response := map[string]interface{}{
				"status": "error",
				"error":  err.Error(),
			}
			respData, _ := json.Marshal(response)
			msg.Respond(respData)
			return
		}

		// Send acknowledgment
		response := map[string]interface{}{
			"status":    "accepted",
			"message":   "VM creation started",
			"requestId": request.RequestID,
			"vmId":      vm.ID,
			"userId":    request.UserID,
		}
		respData, err := json.Marshal(response)
		if err != nil {
			log.Printf("failed to marshal acknowledgment response: %v", err)
		} else {
			if err := msg.Respond(respData); err != nil {
				log.Printf("failed to send acknowledgment response: %v", err)
			} else {
				log.Println("sent acknowledgment response for VM creation request")
			}
		}

		// Start VM creation in a separate goroutine
		go func() {
			ctx := context.Background()

			// Use the requestId if provided, otherwise generate one
			requestID := request.RequestID
			if requestID == "" {
				requestID = fmt.Sprintf("req-%s", GenerateUniqueID(8))
			}

			log.Printf("starting VM creation process (RequestID: %s, UserID: %s)", requestID, request.UserID)

			// Create VM instance on infrastructure
			vm, err := m.createVM(ctx, request.UserID, vmConfig)
			if err != nil {
				log.Printf("failed to create VM (RequestID: %s): %v", requestID, err)

				// Update VM status to failed
				updateErr := m.db.UpdateVMStatus(context.Background(), vm.ID, VMStatusFailed)
				if updateErr != nil {
					log.Printf("failed to update VM status to failed: %v", updateErr)
				}

				// Publish error notification
				errorResp := map[string]interface{}{
					"status":     "error",
					"request_id": requestID,
					"error":      err.Error(),
				}
				errData, _ := json.Marshal(errorResp)
				m.natsClient.Publish("vm.creation.error", errData)
			}

			// Publish created VM info
			vmInfo := &VMCreationResponse{
				VMID:       vm.ID,
				ExoscaleID: vm.ExoscaleID,
				RequestID:  requestID,
				UserID:     request.UserID,
				Name:       vm.Name,
				IP:         vm.IP,
				ExpiresAt:  vm.ExpiresAt.Unix(),
				Status:     vm.Status,
			}

			vmData, err := json.Marshal(vmInfo)
			if err != nil {
				log.Printf("failed to marshal VM data (RequestID: %s): %v", requestID, err)
				return
			}

			if err := m.natsClient.Publish("vm.created", vmData); err != nil {
				log.Printf("failed to publish vm.created message (RequestID: %s): %v", requestID, err)
				return
			}

			log.Printf("published VM on vm.created - Name: %s, ID: %s, Exoscale ID: %s, IP: %s", vm.Name, vm.ID, vm.ExoscaleID, vm.IP)
		}()
	})

	if err != nil {
		return fmt.Errorf("failed to subscribe to vm.create: %w", err)
	}

	// Subscribe to session check queries
	_, err = m.natsClient.Subscribe("vm.check", func(msg *nats.Msg) {
		log.Println("received request on vm.check")

		// Only proceed if we have a Reply subject
		if msg.Reply == "" {
			log.Println("no reply subject for VM session check, ignoring request")
			return
		}

		var request struct {
			UserID     string `json:"userId"`
			TerminalID string `json:"terminalId"`
		}

		if err := json.Unmarshal(msg.Data, &request); err != nil {
			log.Printf("failed to unmarshal VM session check request: %v", err)

			// Reply with error response
			errResponse := map[string]string{
				"status": "error",
				"error":  "invalid request format",
			}
			errData, _ := json.Marshal(errResponse)
			m.natsClient.Publish(msg.Reply, errData)
			return
		}

		// Process the request in a separate goroutine
		go func() {
			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			defer cancel()

			// Check if the user is in cooldown
			inCooldown, cooldownInfo, err := m.db.CheckUserCooldown(ctx, request.UserID)
			if err != nil {
				log.Printf("error checking cooldown for user %s: %v", request.UserID, err)
			}
			if inCooldown {
				log.Printf("user %s is currently in cooldown", request.UserID)
				response := map[string]interface{}{
					"status":   "cooldown",
					"cooldown": cooldownInfo,
				}
				respData, _ := json.Marshal(response)
				m.natsClient.Publish(msg.Reply, respData)
				return
			} else {
				log.Printf("user %s is not in cooldown", request.UserID)
			}

			// Check if user has an existing VM (either in "active" or "creating" status)
			vm, err := m.db.GetUserVM(ctx, request.UserID)
			if err != nil {
				log.Printf("error retrieving VM for user %s: %v", request.UserID, err)
				errResp := map[string]interface{}{
					"status": "error",
					"error":  err.Error(),
				}
				errData, _ := json.Marshal(errResp)
				m.natsClient.Publish(msg.Reply, errData)
				return
			}

			// Check status of existing VM is one is found
			if vm != nil {
				log.Printf("[%s] VM found for user %s", vm.Status, request.UserID)

				var response map[string]interface{}
				response = make(map[string]interface{})

				// Handle VMStatusCreating status
				if vm.Status == VMStatusCreating {
					response["userId"] = request.UserID
					response["vmId"] = vm.ID
					response["message"] = "VM under creation found"
					response["status"] = vm.Status
				}

				// Handle VMStatusActive status
				if vm.Status == VMStatusActive {
					response["userId"] = request.UserID
					response["vmId"] = vm.ID
					response["vmIp"] = vm.IP
					response["message"] = "active VM found"
					response["status"] = vm.Status
				}

				respData, _ := json.Marshal(response)
				m.natsClient.Publish(msg.Reply, respData)
				return
			}

			// Need to create a VM as none was found for the user
			response := map[string]interface{}{
				"status": "need_vm",
			}
			respData, _ := json.Marshal(response)
			m.natsClient.Publish(msg.Reply, respData)
		}()
	})

	if err != nil {
		return fmt.Errorf("failed to subscribe to vm.check: %w", err)
	}

	log.Println("NATS subscriptions set up successfully")
	return nil
}

// Checks for expired VMs and initiates deletion process
func (m *VMManager) deleteExpiredVMs(ctx context.Context) error {
	log.Printf("deleting expired VMs...")

	expiredVMs, err := m.db.GetExpiredVMs(ctx)
	if err != nil {
		return fmt.Errorf("error getting expired VMs: %w", err)
	}

	// Update VM status
	for _, vm := range expiredVMs {
		if err := m.db.UpdateVMStatus(ctx, vm.ID, VMStatusPendingDeletion); err != nil {
			log.Printf("error updating VM %s status: %v", vm.ID, err)
			continue
		}

		// Define cooldown expiration
		expiresAt := time.Now().Add(m.config.CooldownExpirationTime)
		formattedTime := fmt.Sprintf("%s (GMT)", expiresAt.Format("15:04"))

		// Create a cooldown entry for user
		if err := m.db.SetUserCooldown(ctx, vm.UserID, expiresAt, "session_timeout"); err != nil {
			log.Printf("error setting user %s in cooldown: %v", vm.UserID, err)
			continue
		}

		// Create expiration message
		notification := map[string]interface{}{
			"userId": vm.UserID,
			"vmId":   vm.ID,
			"reason": "session_timeout",
			"cooldown": map[string]interface{}{
				"expiryTimestamp": expiresAt.UnixMilli(),
				"formattedTime":   formattedTime,
			},
		}

		notifData, err := json.Marshal(notification)
		if err != nil {
			log.Printf("error marshaling VM expiration notification: %v", err)
			continue
		}

		// Send Request and wait for Reply with timeout
		msg, err := m.natsClient.Request("vm.expiring", notifData, 5*time.Second)
		if err != nil {
			log.Printf("error or timeout on VM expiration notification for VM %s: %v", vm.ID, err)
			// Handle error case - maybe retry or mark VM for forced deletion
			continue
		}

		// Process response
		var response struct {
			Acknowledged bool   `json:"acknowledged"`
			Error        string `json:"error,omitempty"`
		}

		if err := json.Unmarshal(msg.Data, &response); err != nil {
			log.Printf("error parsing expiration response for VM %s: %v", vm.ID, err)
			continue
		}

		if response.Acknowledged && response.Error == "" {
			log.Printf("successfully notified clients of VM %s expiration, proceeding with deletion", vm.ID)

			// Delete VM instance
			if err := m.deleteVM(ctx, vm.ExoscaleID); err != nil {
				log.Printf("error deleting VM %s resources: %v", vm.ID, err)
				continue
			}

			// Mark VM as deleted in database
			if err := m.db.MarkVMDeleted(ctx, vm.ID); err != nil {
				log.Printf("error marking VM as deleted in database: %v", err)
			}
		} else {
			log.Printf("failed to notify clients of VM %s expiration: %s", vm.ID, response.Error)
		}
	}

	return nil
}

// Starts a goroutine to periodically check for VM expirations
func (m *VMManager) StartExpirationChecker(ctx context.Context, checkInterval time.Duration) {
	log.Printf("starting VM expiration checker with interval %s", checkInterval)

	ticker := time.NewTicker(checkInterval)
	go func() {
		for {
			select {
			case <-ctx.Done():
				ticker.Stop()
				log.Println("VM expiration checker stopped")
				return
			case <-ticker.C:
				if err := m.deleteExpiredVMs(ctx); err != nil {
					log.Printf("Error checking expired VMs: %v", err)
				}
			}
		}
	}()
}

// Starts a goroutine to periodically remove expired cooldowns
func (m *VMManager) StartCooldownExpirationChecker(ctx context.Context, checkInterval time.Duration) {
	log.Printf("starting cooldown expiration checker with interval %s", checkInterval)

	ticker := time.NewTicker(checkInterval)
	go func() {
		for {
			select {
			case <-ctx.Done():
				ticker.Stop()
				log.Println("cooldown expiration checker stopped")
				return
			case <-ticker.C:
				if err := m.db.deleteExpiredCooldowns(ctx); err != nil {
					log.Printf("error checking expired Cooldowns: %v", err)
				}
			}
		}
	}()
}