rmm2/backend/api/scheduler.go

package api

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"

	"github.com/cynfo/rmm-backend/db"
	"github.com/cynfo/rmm-backend/models"
	"github.com/cynfo/rmm-backend/ws"
)

type TaskScheduler struct {
	db  *db.Database
	hub *ws.Hub
}

func NewTaskScheduler(database *db.Database, hub *ws.Hub) *TaskScheduler {
	return &TaskScheduler{db: database, hub: hub}
}

func (s *TaskScheduler) Run() {
	log.Println("Task-Scheduler gestartet (30s Intervall)")
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	// Run once immediately
	s.processDueTasks()

	for range ticker.C {
		s.processDueTasks()
	}
}

func (s *TaskScheduler) processDueTasks() {
	tasks, err := s.db.GetDueTasks()
	if err != nil {
		log.Printf("Scheduler: Fehler beim Laden faelliger Tasks: %v", err)
		return
	}

	for _, task := range tasks {
		go s.executeTask(task)
	}
}

func (s *TaskScheduler) executeTask(task models.Task) {
	log.Printf("Scheduler: Starte Task #%d (%s) fuer Agent %s", task.ID, task.Action, task.AgentID)

	// Set status to running
	s.db.UpdateTaskStatus(task.ID, "running", nil)

	var result interface{}
	var taskErr error

	switch task.Action {
	case "backup":
		result, taskErr = s.executeBackup(task.AgentID)
	case "update":
		result, taskErr = s.executeUpdateWithHealthCheck(task.AgentID, task.Reboot, task.HealthCheck, task.HealthCheckTimeout)
	default:
		taskErr = fmt.Errorf("unbekannte Aktion: %s", task.Action)
	}

	success := taskErr == nil
	healthCheckWarning := false
	if !success {
		// Check if it's a health check warning (update ran OK but agent didn't come back)
		if resultMap, ok := result.(map[string]interface{}); ok {
			if hc, ok := resultMap["health_check"].(map[string]interface{}); ok {
				if online, ok := hc["agent_online"].(bool); ok && !online {
					healthCheckWarning = true
				}
			}
		}
		if !healthCheckWarning {
			result = map[string]string{"error": taskErr.Error()}
		}
		log.Printf("Scheduler: Task #%d fehlgeschlagen: %v", task.ID, taskErr)
	} else {
		log.Printf("Scheduler: Task #%d erfolgreich abgeschlossen", task.ID)
	}

	// Use CompleteRecurringTask for proper recurring handling
	// Health check warnings count as success for recurring schedule, but status = "completed_with_warning"
	scheduleSuccess := success || healthCheckWarning
	s.db.CompleteRecurringTask(task.ID, task.ScheduleType, task.ScheduleTime,
		task.ScheduleWeekday, task.ScheduleMonthday, scheduleSuccess, result)

	// Override status for health check warnings
	if healthCheckWarning {
		s.db.UpdateTaskStatusDirect(task.ID, "completed_with_warning")
	}

	// Re-read task for webhook payload
	fullTask, _ := s.db.GetTask(task.ID)

	// Send webhook if configured
	if task.WebhookURL != "" && fullTask != nil {
		s.sendWebhook(task.ID, task.WebhookURL, fullTask)
	}
}

func (s *TaskScheduler) executeBackup(agentID string) (interface{}, error) {
	if !s.hub.IsAgentConnected(agentID) {
		return nil, fmt.Errorf("Agent nicht verbunden")
	}

	cmdID := generateID()
	cmd := map[string]interface{}{
		"type":    "command",
		"id":      cmdID,
		"command": "backup",
		"params":  map[string]interface{}{},
	}

	responseCh := s.hub.RegisterResponseWaiter(cmdID)
	defer s.hub.UnregisterResponseWaiter(cmdID)

	cmdJSON, _ := json.Marshal(cmd)
	if err := s.hub.SendToAgent(agentID, cmdJSON); err != nil {
		return nil, fmt.Errorf("Command senden fehlgeschlagen: %w", err)
	}

	select {
	case resp := <-responseCh:
		respStatus, _ := resp["status"].(string)
		if respStatus != "ok" {
			errMsg, _ := resp["error"].(string)
			return resp, fmt.Errorf("Backup fehlgeschlagen: %s", errMsg)
		}

		data, ok := resp["data"].(map[string]interface{})
		if !ok {
			return resp, fmt.Errorf("ungueltiges Response-Format")
		}

		configB64, _ := data["config"].(string)
		hash, _ := data["hash"].(string)
		sizeFloat, _ := data["size"].(float64)
		size := int(sizeFloat)

		configXML, err := base64.StdEncoding.DecodeString(configB64)
		if err != nil {
			return nil, fmt.Errorf("Base64-Dekodierung fehlgeschlagen")
		}

		backupID, isNew, err := s.db.SaveConfigBackup(agentID, hash, size, string(configXML))
		if err != nil {
			return nil, fmt.Errorf("Backup speichern fehlgeschlagen: %w", err)
		}

		msg := "Backup gespeichert"
		if !isNew {
			msg = "Config unveraendert (Hash identisch)"
		}

		return map[string]interface{}{
			"message":   msg,
			"backup_id": backupID,
			"sha256":    hash,
			"size":      size,
			"new":       isNew,
		}, nil

	case <-time.After(60 * time.Second):
		return nil, fmt.Errorf("Timeout: Agent hat nicht rechtzeitig geantwortet")
	}
}

func (s *TaskScheduler) executeUpdate(agentID string, reboot bool) (interface{}, error) {
	return s.executeUpdateWithHealthCheck(agentID, reboot, false, 600)
}

func (s *TaskScheduler) executeUpdateWithHealthCheck(agentID string, reboot bool, healthCheck bool, healthCheckTimeout int) (interface{}, error) {
	if !s.hub.IsAgentConnected(agentID) {
		return nil, fmt.Errorf("Agent nicht verbunden")
	}

	// First: update_check
	checkCmdID := generateID()
	checkCmd := map[string]interface{}{
		"type":    "command",
		"id":      checkCmdID,
		"command": "update_check",
	}

	checkCh := s.hub.RegisterResponseWaiter(checkCmdID)
	defer s.hub.UnregisterResponseWaiter(checkCmdID)

	checkJSON, _ := json.Marshal(checkCmd)
	if err := s.hub.SendToAgent(agentID, checkJSON); err != nil {
		return nil, fmt.Errorf("Update-Check senden fehlgeschlagen: %w", err)
	}

	var checkResult map[string]interface{}
	select {
	case resp := <-checkCh:
		checkResult = resp
	case <-time.After(120 * time.Second):
		return nil, fmt.Errorf("Timeout beim Update-Check")
	}

	// Then: update
	updateCmdID := generateID()
	updateCmd := map[string]interface{}{
		"type":    "command",
		"id":      updateCmdID,
		"command": "update",
		"params":  map[string]interface{}{"reboot": reboot},
	}

	updateCh := s.hub.RegisterResponseWaiter(updateCmdID)
	defer s.hub.UnregisterResponseWaiter(updateCmdID)

	updateJSON, _ := json.Marshal(updateCmd)
	if err := s.hub.SendToAgent(agentID, updateJSON); err != nil {
		return nil, fmt.Errorf("Update-Command senden fehlgeschlagen: %w", err)
	}

	var updateResult map[string]interface{}
	select {
	case resp := <-updateCh:
		respStatus, _ := resp["status"].(string)
		if respStatus != "ok" {
			errMsg, _ := resp["error"].(string)
			return resp, fmt.Errorf("Update fehlgeschlagen: %s", errMsg)
		}
		updateResult = resp
	case <-time.After(900 * time.Second):
		return nil, fmt.Errorf("Timeout beim Update")
	}

	result := map[string]interface{}{
		"message":       "Update ausgefuehrt",
		"check_result":  checkResult,
		"update_result": updateResult,
		"reboot":        reboot,
	}

	// Health Check: wait for agent to come back online
	if healthCheck {
		timeout := time.Duration(healthCheckTimeout) * time.Second
		if timeout <= 0 {
			timeout = 600 * time.Second
		}
		log.Printf("Scheduler: Health Check fuer Agent %s (Timeout: %v)", agentID, timeout)

		// If no reboot, agent should still be connected — give it a moment for post-update restart
		if !reboot {
			time.Sleep(15 * time.Second)
		}

		checkedAt := time.Now().UTC().Format(time.RFC3339)
		agentOnline := s.hub.WaitForAgentOnline(agentID, timeout)
		result["health_check"] = map[string]interface{}{
			"agent_online": agentOnline,
			"checked_at":   checkedAt,
		}

		if !agentOnline {
			result["message"] = "Update ausgefuehrt, aber Agent nicht wieder online (Health Check fehlgeschlagen)"
			log.Printf("Scheduler: Health Check fehlgeschlagen — Agent %s nach %v nicht online", agentID, timeout)
			// Return result but signal warning via error
			return result, fmt.Errorf("Health Check: Agent nach %v nicht online", timeout)
		}
		log.Printf("Scheduler: Health Check OK — Agent %s wieder online", agentID)
	}

	return result, nil
}

func (s *TaskScheduler) sendWebhook(taskID int, url string, task interface{}) {
	payload, err := json.Marshal(task)
	if err != nil {
		log.Printf("Scheduler: Webhook-Payload fuer Task #%d serialisieren fehlgeschlagen: %v", taskID, err)
		s.db.SetTaskWebhookResult(taskID, false, fmt.Sprintf("JSON-Fehler: %v", err))
		return
	}

	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Post(url, "application/json", bytes.NewReader(payload))
	if err != nil {
		log.Printf("Scheduler: Webhook fuer Task #%d fehlgeschlagen: %v", taskID, err)
		s.db.SetTaskWebhookResult(taskID, false, fmt.Sprintf("Fehler: %v", err))
		return
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
	response := fmt.Sprintf("HTTP %d: %s", resp.StatusCode, string(body))

	sent := resp.StatusCode >= 200 && resp.StatusCode < 300
	s.db.SetTaskWebhookResult(taskID, sent, response)
	log.Printf("Scheduler: Webhook fuer Task #%d gesendet (Status: %d)", taskID, resp.StatusCode)
}