Redflag/aggregator-server/internal/database/queries/subsystems.go

package queries

import (
	"database/sql"
	"fmt"

	"github.com/Fimeg/RedFlag/aggregator-server/internal/models"
	"github.com/google/uuid"
	"github.com/jmoiron/sqlx"
)

type SubsystemQueries struct {
	db *sqlx.DB
}

func NewSubsystemQueries(db *sqlx.DB) *SubsystemQueries {
	return &SubsystemQueries{db: db}
}

// GetSubsystems retrieves all subsystems for an agent
func (q *SubsystemQueries) GetSubsystems(agentID uuid.UUID) ([]models.AgentSubsystem, error) {
	query := `
		SELECT id, agent_id, subsystem, enabled, interval_minutes, auto_run,
		       last_run_at, next_run_at, created_at, updated_at
		FROM agent_subsystems
		WHERE agent_id = $1
		ORDER BY subsystem
	`

	var subsystems []models.AgentSubsystem
	err := q.db.Select(&subsystems, query, agentID)
	if err != nil {
		return nil, fmt.Errorf("failed to get subsystems: %w", err)
	}

	return subsystems, nil
}

// GetSubsystem retrieves a specific subsystem for an agent
func (q *SubsystemQueries) GetSubsystem(agentID uuid.UUID, subsystem string) (*models.AgentSubsystem, error) {
	query := `
		SELECT id, agent_id, subsystem, enabled, interval_minutes, auto_run,
		       last_run_at, next_run_at, created_at, updated_at
		FROM agent_subsystems
		WHERE agent_id = $1 AND subsystem = $2
	`

	var sub models.AgentSubsystem
	err := q.db.Get(&sub, query, agentID, subsystem)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get subsystem: %w", err)
	}

	return &sub, nil
}

// UpdateSubsystem updates a subsystem configuration
func (q *SubsystemQueries) UpdateSubsystem(agentID uuid.UUID, subsystem string, config models.SubsystemConfig) error {
	// Build dynamic update query based on provided fields
	updates := []string{}
	args := []interface{}{agentID, subsystem}
	argIdx := 3

	if config.Enabled != nil {
		updates = append(updates, fmt.Sprintf("enabled = $%d", argIdx))
		args = append(args, *config.Enabled)
		argIdx++
	}

	if config.IntervalMinutes != nil {
		updates = append(updates, fmt.Sprintf("interval_minutes = $%d", argIdx))
		args = append(args, *config.IntervalMinutes)
		argIdx++
	}

	if config.AutoRun != nil {
		updates = append(updates, fmt.Sprintf("auto_run = $%d", argIdx))
		args = append(args, *config.AutoRun)
		argIdx++

		// If enabling auto_run, calculate next_run_at
		if *config.AutoRun {
			updates = append(updates, fmt.Sprintf("next_run_at = NOW() + INTERVAL '%d minutes'", argIdx))
		}
	}

	if len(updates) == 0 {
		return fmt.Errorf("no fields to update")
	}

	updates = append(updates, "updated_at = NOW()")

	query := fmt.Sprintf(`
		UPDATE agent_subsystems
		SET %s
		WHERE agent_id = $1 AND subsystem = $2
	`, joinUpdates(updates))

	result, err := q.db.Exec(query, args...)
	if err != nil {
		return fmt.Errorf("failed to update subsystem: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rows == 0 {
		return fmt.Errorf("subsystem not found")
	}

	return nil
}

// UpdateLastRun updates the last_run_at timestamp for a subsystem
func (q *SubsystemQueries) UpdateLastRun(agentID uuid.UUID, subsystem string) error {
	query := `
		UPDATE agent_subsystems
		SET last_run_at = NOW(),
		    next_run_at = CASE
		        WHEN auto_run THEN NOW() + (interval_minutes || ' minutes')::INTERVAL
		        ELSE next_run_at
		    END,
		    updated_at = NOW()
		WHERE agent_id = $1 AND subsystem = $2
	`

	result, err := q.db.Exec(query, agentID, subsystem)
	if err != nil {
		return fmt.Errorf("failed to update last run: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rows == 0 {
		return fmt.Errorf("subsystem not found")
	}

	return nil
}

// GetDueSubsystems retrieves all subsystems that are due to run
func (q *SubsystemQueries) GetDueSubsystems() ([]models.AgentSubsystem, error) {
	query := `
		SELECT id, agent_id, subsystem, enabled, interval_minutes, auto_run,
		       last_run_at, next_run_at, created_at, updated_at
		FROM agent_subsystems
		WHERE enabled = true
		  AND auto_run = true
		  AND (next_run_at IS NULL OR next_run_at <= NOW())
		ORDER BY next_run_at ASC NULLS FIRST
		LIMIT 1000
	`

	var subsystems []models.AgentSubsystem
	err := q.db.Select(&subsystems, query)
	if err != nil {
		return nil, fmt.Errorf("failed to get due subsystems: %w", err)
	}

	return subsystems, nil
}

// GetSubsystemStats retrieves statistics for a subsystem
func (q *SubsystemQueries) GetSubsystemStats(agentID uuid.UUID, subsystem string) (*models.SubsystemStats, error) {
	query := `
		SELECT
		    s.subsystem,
		    s.enabled,
		    s.last_run_at,
		    s.next_run_at,
		    s.interval_minutes,
		    s.auto_run,
		    COUNT(c.id) FILTER (WHERE c.command_type = 'scan_' || s.subsystem) as run_count,
		    MAX(c.status) FILTER (WHERE c.command_type = 'scan_' || s.subsystem) as last_status,
		    MAX(al.duration_seconds) FILTER (WHERE al.action = 'scan_' || s.subsystem) as last_duration
		FROM agent_subsystems s
		LEFT JOIN agent_commands c ON c.agent_id = s.agent_id
		LEFT JOIN agent_logs al ON al.command_id = c.id
		WHERE s.agent_id = $1 AND s.subsystem = $2
		GROUP BY s.subsystem, s.enabled, s.last_run_at, s.next_run_at, s.interval_minutes, s.auto_run
	`

	var stats models.SubsystemStats
	err := q.db.Get(&stats, query, agentID, subsystem)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get subsystem stats: %w", err)
	}

	return &stats, nil
}

// EnableSubsystem enables a subsystem
func (q *SubsystemQueries) EnableSubsystem(agentID uuid.UUID, subsystem string) error {
	enabled := true
	return q.UpdateSubsystem(agentID, subsystem, models.SubsystemConfig{
		Enabled: &enabled,
	})
}

// DisableSubsystem disables a subsystem
func (q *SubsystemQueries) DisableSubsystem(agentID uuid.UUID, subsystem string) error {
	enabled := false
	return q.UpdateSubsystem(agentID, subsystem, models.SubsystemConfig{
		Enabled: &enabled,
	})
}

// SetAutoRun enables or disables auto-run for a subsystem
func (q *SubsystemQueries) SetAutoRun(agentID uuid.UUID, subsystem string, autoRun bool) error {
	return q.UpdateSubsystem(agentID, subsystem, models.SubsystemConfig{
		AutoRun: &autoRun,
	})
}

// SetInterval sets the interval for a subsystem
func (q *SubsystemQueries) SetInterval(agentID uuid.UUID, subsystem string, intervalMinutes int) error {
	return q.UpdateSubsystem(agentID, subsystem, models.SubsystemConfig{
		IntervalMinutes: &intervalMinutes,
	})
}

// CreateSubsystem creates a new subsystem configuration (used for custom subsystems)
func (q *SubsystemQueries) CreateSubsystem(sub *models.AgentSubsystem) error {
	query := `
		INSERT INTO agent_subsystems (agent_id, subsystem, enabled, interval_minutes, auto_run, last_run_at, next_run_at)
		VALUES ($1, $2, $3, $4, $5, $6, $7)
		RETURNING id, created_at, updated_at
	`

	err := q.db.QueryRow(
		query,
		sub.AgentID,
		sub.Subsystem,
		sub.Enabled,
		sub.IntervalMinutes,
		sub.AutoRun,
		sub.LastRunAt,
		sub.NextRunAt,
	).Scan(&sub.ID, &sub.CreatedAt, &sub.UpdatedAt)

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

	return nil
}

// DeleteSubsystem deletes a subsystem configuration
func (q *SubsystemQueries) DeleteSubsystem(agentID uuid.UUID, subsystem string) error {
	query := `
		DELETE FROM agent_subsystems
		WHERE agent_id = $1 AND subsystem = $2
	`

	result, err := q.db.Exec(query, agentID, subsystem)
	if err != nil {
		return fmt.Errorf("failed to delete subsystem: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rows == 0 {
		return fmt.Errorf("subsystem not found")
	}

	return nil
}

// Helper function to join update statements
func joinUpdates(updates []string) string {
	result := ""
	for i, update := range updates {
		if i > 0 {
			result += ", "
		}
		result += update
	}
	return result
}