RedFlag/server/internal/services/timeout.go

package services

import (
	"fmt"
	"log"
	"time"

	"github.com/Fimeg/RedFlag/server/internal/database/queries"
	"github.com/Fimeg/RedFlag/server/internal/models"
	"github.com/google/uuid"
)

// TimeoutService handles timeout management for long-running operations
type TimeoutService struct {
	commandQueries  *queries.CommandQueries
	updateQueries   *queries.UpdateQueries
	agentQueries    *queries.AgentQueries
	ticker          *time.Ticker
	stopChan        chan bool
	sentTimeout     time.Duration // For commands already sent to agents
	pendingTimeout  time.Duration // For commands stuck in queue
	receivedTimeout time.Duration // For commands received by agent but not completed (Migration 033 §4)
	updateTimeout   time.Duration // For agents stuck in is_updating=true
	checkInterval   time.Duration // How often to check for timeouts
}

// NewTimeoutService creates a new timeout service with configurable durations.
// Pass zero values to use defaults (2h sent, 30m pending, 30m received, 15m update,
// 5m check interval).
func NewTimeoutService(cq *queries.CommandQueries, uq *queries.UpdateQueries, aq *queries.AgentQueries, sentTimeout, pendingTimeout, receivedTimeout, updateTimeout, checkInterval time.Duration) *TimeoutService {
	if sentTimeout <= 0 {
		sentTimeout = 2 * time.Hour
	}
	if pendingTimeout <= 0 {
		pendingTimeout = 30 * time.Minute
	}
	if receivedTimeout <= 0 {
		receivedTimeout = 30 * time.Minute
	}
	if updateTimeout <= 0 {
		updateTimeout = 15 * time.Minute
	}
	if checkInterval <= 0 {
		checkInterval = 5 * time.Minute
	}
	return &TimeoutService{
		commandQueries:  cq,
		updateQueries:   uq,
		agentQueries:    aq,
		sentTimeout:     sentTimeout,
		pendingTimeout:  pendingTimeout,
		receivedTimeout: receivedTimeout,
		updateTimeout:   updateTimeout,
		checkInterval:   checkInterval,
		stopChan:        make(chan bool),
	}
}

// Start begins the timeout monitoring service
func (ts *TimeoutService) Start() {
	log.Printf("[INFO] [server] [timeout] service_started sent_timeout=%v pending_timeout=%v received_timeout=%v update_timeout=%v check_interval=%v",
		ts.sentTimeout, ts.pendingTimeout, ts.receivedTimeout, ts.updateTimeout, ts.checkInterval)

	ts.ticker = time.NewTicker(ts.checkInterval)

	go func() {
		for {
			select {
			case <-ts.ticker.C:
				ts.checkForTimeouts()
				ts.checkForReceivedTimeouts()
				ts.reconcileAgentUpdates()
			case <-ts.stopChan:
				ts.ticker.Stop()
				log.Println("Timeout service stopped")
				return
			}
		}
	}()
}

// Stop stops the timeout monitoring service
func (ts *TimeoutService) Stop() {
	close(ts.stopChan)
}

// checkForTimeouts checks for commands that have been running too long
func (ts *TimeoutService) checkForTimeouts() {
	log.Println("Checking for timed out operations...")

	sentTimeoutThreshold := time.Now().Add(-ts.sentTimeout)
	pendingTimeoutThreshold := time.Now().Add(-ts.pendingTimeout)
	timedOutCommands := make([]models.AgentCommand, 0)

	// Check 'sent' commands (configurable, default 2 hours)
	sentCommands, err := ts.commandQueries.GetCommandsByStatus(models.CommandStatusSent)
	if err != nil {
		log.Printf("Error getting sent commands: %v", err)
	} else {
		for _, command := range sentCommands {
			// Check if command has been sent and is older than sent timeout threshold
			if command.SentAt != nil && command.SentAt.Before(sentTimeoutThreshold) {
				timedOutCommands = append(timedOutCommands, command)
			}
		}
	}

	// Check 'pending' commands (configurable, default 30 minutes)
	pendingCommands, err := ts.commandQueries.GetCommandsByStatus(models.CommandStatusPending)
	if err != nil {
		log.Printf("Error getting pending commands: %v", err)
	} else {
		for _, command := range pendingCommands {
			// Check if command has been pending longer than pending timeout threshold
			if command.CreatedAt.Before(pendingTimeoutThreshold) {
				timedOutCommands = append(timedOutCommands, command)
				log.Printf("Found stuck pending command %s (type: %s, created: %s, age: %v)",
					command.ID, command.CommandType, command.CreatedAt.Format(time.RFC3339), time.Since(command.CreatedAt))
			}
		}
	}

	if len(timedOutCommands) > 0 {
		log.Printf("[INFO] [server] [timeout] timed_out_commands=%d sent_checked=%d pending_checked=%d sent_timeout=%v pending_timeout=%v",
			len(timedOutCommands), len(sentCommands), len(pendingCommands), ts.sentTimeout, ts.pendingTimeout)

		for _, command := range timedOutCommands {
			if err := ts.timeoutCommand(&command); err != nil {
				log.Printf("Error timing out command %s: %v", command.ID, err)
			}
		}
	} else {
		log.Println("No timed out operations found")
	}
}

// timeoutCommand marks a specific command as timed out and updates related entities
func (ts *TimeoutService) timeoutCommand(command *models.AgentCommand) error {
	// Determine which timeout duration was applied
	var appliedTimeout time.Duration
	if command.Status == models.CommandStatusSent {
		appliedTimeout = ts.sentTimeout
	} else {
		appliedTimeout = ts.pendingTimeout
	}

	log.Printf("Timing out command %s (type: %s, agent: %s)",
		command.ID, command.CommandType, command.AgentID)

	// Update command status to timed_out
	if err := ts.commandQueries.UpdateCommandStatus(command.ID, models.CommandStatusTimedOut); err != nil {
		return fmt.Errorf("failed to update command status: %w", err)
	}

	// Update result with timeout information
	result := models.JSONB{
		"error":       "operation timed out",
		"timeout_at":  time.Now(),
		"duration":    appliedTimeout.String(),
		"command_id":  command.ID.String(),
	}

	if err := ts.commandQueries.UpdateCommandResult(command.ID, result); err != nil {
		return fmt.Errorf("failed to update command result: %w", err)
	}

	// Update related update package status if applicable
	if err := ts.updateRelatedPackageStatus(command, appliedTimeout); err != nil {
		log.Printf("Warning: failed to update related package status: %v", err)
		// Don't return error here as the main timeout operation succeeded
	}

	// Create a log entry for the timeout
	logEntry := &models.UpdateLog{
		ID:              uuid.New(),
		AgentID:         command.AgentID,
		UpdatePackageID: ts.extractUpdatePackageID(command),
		Action:          command.CommandType,
		Result:          "failed", // Use 'failed' to comply with database constraint
		Stdout:          "",
		Stderr:          fmt.Sprintf("Command %s timed out after %v (timeout_id: %s)", command.CommandType, appliedTimeout, command.ID),
		ExitCode:        124, // Standard timeout exit code
		DurationSeconds: int(appliedTimeout.Seconds()),
		ExecutedAt:      time.Now(),
	}

	if err := ts.updateQueries.CreateUpdateLog(logEntry); err != nil {
		log.Printf("Warning: failed to create timeout log entry: %v", err)
		// Don't return error here as the main timeout operation succeeded
	}

	log.Printf("Successfully timed out command %s", command.ID)
	return nil
}

// updateRelatedPackageStatus updates the status of related update packages when a command times out
func (ts *TimeoutService) updateRelatedPackageStatus(command *models.AgentCommand, appliedTimeout time.Duration) error {
	// Extract update_id from command params if it exists
	_, ok := command.Params["update_id"].(string)
	if !ok {
		// This command doesn't have an associated update_id, so no package status to update
		return nil
	}

	// Update the package status to 'failed' with timeout reason
	metadata := models.JSONB{
		"timeout":       true,
		"timeout_at":    time.Now(),
		"timeout_duration": appliedTimeout.String(),
		"command_id":    command.ID.String(),
		"failure_reason": "operation timed out",
	}

	return ts.updateQueries.UpdatePackageStatus(command.AgentID,
		command.Params["package_type"].(string),
		command.Params["package_name"].(string),
		"failed",
		metadata,
		nil) // nil = use time.Now() for timeout operations
}

// extractUpdatePackageID extracts the update package ID from command params
func (ts *TimeoutService) extractUpdatePackageID(command *models.AgentCommand) *uuid.UUID {
	updateIDStr, ok := command.Params["update_id"].(string)
	if !ok {
		return nil
	}

	updateID, err := uuid.Parse(updateIDStr)
	if err != nil {
		return nil
	}

	return &updateID
}

// GetTimeoutStatus returns statistics about timed out operations
func (ts *TimeoutService) GetTimeoutStatus() (map[string]interface{}, error) {
	// Get all timed out commands
	timedOutCommands, err := ts.commandQueries.GetCommandsByStatus(models.CommandStatusTimedOut)
	if err != nil {
		return nil, fmt.Errorf("failed to get timed out commands: %w", err)
	}

	// Get all active commands
	activeCommands, err := ts.commandQueries.GetCommandsByStatus(models.CommandStatusSent)
	if err != nil {
		return nil, fmt.Errorf("failed to get active commands: %w", err)
	}

	// Count commands approaching timeout (within 5 minutes of timeout)
	timeoutThreshold := time.Now().Add(-ts.sentTimeout + 5*time.Minute)
	approachingTimeout := 0
	for _, command := range activeCommands {
		if command.SentAt != nil && command.SentAt.Before(timeoutThreshold) {
			approachingTimeout++
		}
	}

	return map[string]interface{}{
		"total_timed_out":          len(timedOutCommands),
		"total_active":             len(activeCommands),
		"approaching_timeout":      approachingTimeout,
		"sent_timeout_duration":    ts.sentTimeout.String(),
		"pending_timeout_duration": ts.pendingTimeout.String(),
		"last_check":               time.Now(),
	}, nil
}

// SetTimeoutDuration allows changing the timeout duration for sent commands
// TODO: This should be deprecated in favor of SetSentTimeout and SetPendingTimeout
func (ts *TimeoutService) SetTimeoutDuration(duration time.Duration) {
	ts.sentTimeout = duration
	log.Printf("Sent timeout duration updated to %v", duration)
}

// SetSentTimeout allows changing the timeout duration for sent commands
func (ts *TimeoutService) SetSentTimeout(duration time.Duration) {
	ts.sentTimeout = duration
	log.Printf("Sent timeout duration updated to %v", duration)
}

// SetPendingTimeout allows changing the timeout duration for pending commands
func (ts *TimeoutService) SetPendingTimeout(duration time.Duration) {
	ts.pendingTimeout = duration
	log.Printf("Pending timeout duration updated to %v", duration)
}

// checkForReceivedTimeouts handles commands the agent received but never completed.
// Distinct from sent-timeouts because we know the agent had it — re-issuance won't
// help; the right action is to mark timed_out so the operator (or scheduler) can
// decide whether to retry or escalate.
//
// Doctrine: TODO-full-command-lifecycle.md §4. Migration 033 added the 'received' state.
func (ts *TimeoutService) checkForReceivedTimeouts() {
	tx, err := ts.commandQueries.DB().Beginx()
	if err != nil {
		log.Printf("[ERROR] [server] [timeout] received_tx_begin_failed error=%v", err)
		return
	}
	defer tx.Rollback()

	stuck, err := ts.commandQueries.GetStuckReceivedCommandsTx(tx, ts.receivedTimeout)
	if err != nil {
		log.Printf("[ERROR] [server] [timeout] get_stuck_received_failed error=%v", err)
		return
	}
	if len(stuck) == 0 {
		return
	}

	for _, command := range stuck {
		cmd := command
		if err := ts.timeoutCommand(&cmd); err != nil {
			log.Printf("[ERROR] [server] [timeout] timeout_received_failed command_id=%s error=%v", cmd.ID, err)
		}
	}

	if err := tx.Commit(); err != nil {
		log.Printf("[ERROR] [server] [timeout] received_tx_commit_failed error=%v", err)
		return
	}

	log.Printf("[INFO] [server] [timeout] received_commands_timed_out count=%d threshold=%v", len(stuck), ts.receivedTimeout)
}

// reconcileAgentUpdates handles agents whose is_updating flag was set but never cleared.
// Decision per agent:
//   - if current_version == updating_to_version: the agent completed the update and reported
//     its new version, but the per-request side-effect to clear is_updating was deliberately
//     omitted. Clear it now (success path).
//   - otherwise: the update never reported in (binary failed to start, network dead, etc).
//     Clear is_updating and log a timeout system_event so the operator sees what happened.
//
// Doctrine: TODO-full-command-lifecycle.md §4 + audit finding "completion loop has no firing path".
func (ts *TimeoutService) reconcileAgentUpdates() {
	if ts.agentQueries == nil {
		return // not wired (test path)
	}
	threshold := time.Now().Add(-ts.updateTimeout)
	stuck, err := ts.agentQueries.GetAgentsStuckUpdating(threshold)
	if err != nil {
		log.Printf("[ERROR] [server] [timeout] get_stuck_updating_agents_failed error=%v", err)
		return
	}
	if len(stuck) == 0 {
		return
	}

	successes, timeouts := 0, 0
	for _, agent := range stuck {
		target := ""
		if agent.UpdatingToVersion != nil {
			target = *agent.UpdatingToVersion
		}

		if target != "" && agent.CurrentVersion == target {
			// Success — agent reported the new version; just close the flag.
			if err := ts.agentQueries.CompleteAgentUpdate(agent.ID.String(), agent.CurrentVersion); err != nil {
				log.Printf("[ERROR] [server] [timeout] complete_update_failed agent_id=%s error=%v", agent.ID, err)
				continue
			}
			ts.recordUpdateEvent(agent.ID, "succeeded", "info",
				fmt.Sprintf("Agent update succeeded: %s reported", agent.CurrentVersion),
				map[string]interface{}{"new_version": agent.CurrentVersion, "reconciled_by": "timeout_service"})
			successes++
			continue
		}

		// Timeout — version never matched. Clear the flag so the operator can retry.
		if err := ts.agentQueries.ClearAgentUpdating(agent.ID); err != nil {
			log.Printf("[ERROR] [server] [timeout] clear_updating_failed agent_id=%s error=%v", agent.ID, err)
			continue
		}
		ageMessage := fmt.Sprintf("Agent update timed out after %v without version attestation (current=%s, target=%s)",
			ts.updateTimeout, agent.CurrentVersion, target)
		ts.recordUpdateEvent(agent.ID, "timed_out", "warning", ageMessage,
			map[string]interface{}{
				"current_version": agent.CurrentVersion,
				"target_version":  target,
				"threshold":       ts.updateTimeout.String(),
				"reconciled_by":   "timeout_service",
			})
		timeouts++
	}

	log.Printf("[INFO] [server] [timeout] agent_updates_reconciled stuck=%d succeeded=%d timed_out=%d threshold=%v",
		len(stuck), successes, timeouts, ts.updateTimeout)
}

// recordUpdateEvent writes a system_event row for an update lifecycle transition. Best
// effort — a failure to log shouldn't block the reconciler from continuing on the next agent.
func (ts *TimeoutService) recordUpdateEvent(agentID uuid.UUID, subtype, severity, message string, metadata map[string]interface{}) {
	event := &models.SystemEvent{
		ID:           uuid.New(),
		AgentID:      &agentID,
		EventType:    "agent_update",
		EventSubtype: subtype,
		Severity:     severity,
		Component:    "agent",
		Message:      message,
		Metadata:     metadata,
		CreatedAt:    time.Now(),
	}
	if err := ts.agentQueries.CreateSystemEvent(event); err != nil {
		log.Printf("[WARNING] [server] [timeout] system_event_write_failed agent_id=%s subtype=%s error=%v",
			agentID, subtype, err)
	}
}