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

package queries

import (
	"database/sql"
	"fmt"

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

// DockerQueries handles database operations for Docker images
type DockerQueries struct {
	db *sql.DB
}

func NewDockerQueries(db *sql.DB) *DockerQueries {
	return &DockerQueries{db: db}
}

// CreateDockerEventsBatch creates multiple Docker image events in a single transaction
func (q *DockerQueries) CreateDockerEventsBatch(events []models.StoredDockerImage) error {
	if len(events) == 0 {
		return nil
	}

	tx, err := q.db.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	// Prepare the insert statement
	stmt, err := tx.Prepare(`
		INSERT INTO docker_images (
			id, agent_id, package_type, package_name, current_version, available_version,
			severity, repository_source, metadata, event_type, created_at
		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
		ON CONFLICT (agent_id, package_name, package_type, created_at) DO NOTHING
	`)
	if err != nil {
		return fmt.Errorf("failed to prepare statement: %w", err)
	}
	defer stmt.Close()

	// Insert each event with error isolation
	for _, event := range events {
		_, err := stmt.Exec(
			event.ID,
			event.AgentID,
			event.PackageType,
			event.PackageName,
			event.CurrentVersion,
			event.AvailableVersion,
			event.Severity,
			event.RepositorySource,
			event.Metadata,
			event.EventType,
			event.CreatedAt,
		)
		if err != nil {
			// Log error but continue with other events
			fmt.Printf("Warning: Failed to insert docker image event %s: %v\n", event.ID, err)
			continue
		}
	}

	return tx.Commit()
}

// GetDockerImages retrieves Docker images based on filter criteria
func (q *DockerQueries) GetDockerImages(filter *models.DockerFilter) (*models.DockerResult, error) {
	query := `
		SELECT id, agent_id, package_type, package_name, current_version, available_version,
		       severity, repository_source, metadata, event_type, created_at
		FROM docker_images
		WHERE 1=1
	`
	args := []interface{}{}
	argIndex := 1

	// Build WHERE clause
	if filter.AgentID != nil {
		query += fmt.Sprintf(" AND agent_id = $%d", argIndex)
		args = append(args, *filter.AgentID)
		argIndex++
	}

	if filter.ImageName != nil {
		query += fmt.Sprintf(" AND package_name ILIKE $%d", argIndex)
		args = append(args, "%"+*filter.ImageName+"%")
		argIndex++
	}

	if filter.Registry != nil {
		query += fmt.Sprintf(" AND repository_source ILIKE $%d", argIndex)
		args = append(args, "%"+*filter.Registry+"%")
		argIndex++
	}

	if filter.Severity != nil {
		query += fmt.Sprintf(" AND severity = $%d", argIndex)
		args = append(args, *filter.Severity)
		argIndex++
	}

	if filter.HasUpdates != nil {
		if *filter.HasUpdates {
			query += fmt.Sprintf(" AND current_version != available_version", argIndex)
		} else {
			query += fmt.Sprintf(" AND current_version = available_version", argIndex)
		}
		argIndex++
	}

	// Add ordering and pagination
	query += " ORDER BY created_at DESC"

	if filter.Limit != nil {
		query += fmt.Sprintf(" LIMIT $%d", argIndex)
		args = append(args, *filter.Limit)
		argIndex++
	}

	if filter.Offset != nil {
		query += fmt.Sprintf(" OFFSET $%d", argIndex)
		args = append(args, *filter.Offset)
		argIndex++
	}

	rows, err := q.db.Query(query, args...)
	if err != nil {
		return nil, fmt.Errorf("failed to query docker images: %w", err)
	}
	defer rows.Close()

	var images []models.StoredDockerImage
	for rows.Next() {
		var image models.StoredDockerImage
		err := rows.Scan(
			&image.ID,
			&image.AgentID,
			&image.PackageType,
			&image.PackageName,
			&image.CurrentVersion,
			&image.AvailableVersion,
			&image.Severity,
			&image.RepositorySource,
			&image.Metadata,
			&image.EventType,
			&image.CreatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan docker image: %w", err)
		}
		images = append(images, image)
	}

	// Get total count
	countQuery := `SELECT COUNT(*) FROM docker_images WHERE 1=1`
	countArgs := []interface{}{}
	countIndex := 1

	if filter.AgentID != nil {
		countQuery += fmt.Sprintf(" AND agent_id = $%d", countIndex)
		countArgs = append(countArgs, *filter.AgentID)
		countIndex++
	}

	if filter.ImageName != nil {
		countQuery += fmt.Sprintf(" AND package_name ILIKE $%d", countIndex)
		countArgs = append(countArgs, "%"+*filter.ImageName+"%")
		countIndex++
	}

	if filter.Registry != nil {
		countQuery += fmt.Sprintf(" AND repository_source ILIKE $%d", countIndex)
		countArgs = append(countArgs, "%"+*filter.Registry+"%")
		countIndex++
	}

	if filter.Severity != nil {
		countQuery += fmt.Sprintf(" AND severity = $%d", countIndex)
		countArgs = append(countArgs, *filter.Severity)
		countIndex++
	}

	if filter.HasUpdates != nil {
		if *filter.HasUpdates {
			countQuery += fmt.Sprintf(" AND current_version != available_version", countIndex)
		} else {
			countQuery += fmt.Sprintf(" AND current_version = available_version", countIndex)
		}
		countIndex++
	}

	var total int
	err = q.db.QueryRow(countQuery, countArgs...).Scan(&total)
	if err != nil {
		return nil, fmt.Errorf("failed to count docker images: %w", err)
	}

	// Calculate pagination
	page := 1
	perPage := 50
	if filter.Offset != nil && filter.Limit != nil {
		page = (*filter.Offset / *filter.Limit) + 1
		perPage = *filter.Limit
	}

	return &models.DockerResult{
		Images:  images,
		Total:   total,
		Page:    page,
		PerPage: perPage,
	}, nil
}

// GetDockerImagesByAgentID retrieves Docker images for a specific agent
func (q *DockerQueries) GetDockerImagesByAgentID(agentID uuid.UUID, limit int) ([]models.StoredDockerImage, error) {
	query := `
		SELECT id, agent_id, package_type, package_name, current_version, available_version,
		       severity, repository_source, metadata, event_type, created_at
		FROM docker_images
		WHERE agent_id = $1
		ORDER BY created_at DESC
		LIMIT $2
	`

	rows, err := q.db.Query(query, agentID, limit)
	if err != nil {
		return nil, fmt.Errorf("failed to query docker images by agent: %w", err)
	}
	defer rows.Close()

	var images []models.StoredDockerImage
	for rows.Next() {
		var image models.StoredDockerImage
		err := rows.Scan(
			&image.ID,
			&image.AgentID,
			&image.PackageType,
			&image.PackageName,
			&image.CurrentVersion,
			&image.AvailableVersion,
			&image.Severity,
			&image.RepositorySource,
			&image.Metadata,
			&image.EventType,
			&image.CreatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan docker image: %w", err)
		}
		images = append(images, image)
	}

	return images, nil
}

// GetDockerImagesWithUpdates retrieves Docker images that have available updates
func (q *DockerQueries) GetDockerImagesWithUpdates(limit int) ([]models.StoredDockerImage, error) {
	query := `
		SELECT id, agent_id, package_type, package_name, current_version, available_version,
		       severity, repository_source, metadata, event_type, created_at
		FROM docker_images
		WHERE current_version != available_version
		ORDER BY created_at DESC
		LIMIT $1
	`

	rows, err := q.db.Query(query, limit)
	if err != nil {
		return nil, fmt.Errorf("failed to query docker images with updates: %w", err)
	}
	defer rows.Close()

	var images []models.StoredDockerImage
	for rows.Next() {
		var image models.StoredDockerImage
		err := rows.Scan(
			&image.ID,
			&image.AgentID,
			&image.PackageType,
			&image.PackageName,
			&image.CurrentVersion,
			&image.AvailableVersion,
			&image.Severity,
			&image.RepositorySource,
			&image.Metadata,
			&image.EventType,
			&image.CreatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan docker image: %w", err)
		}
		images = append(images, image)
	}

	return images, nil
}

// DeleteOldDockerImages deletes Docker images older than the specified number of days
func (q *DockerQueries) DeleteOldDockerImages(days int) error {
	query := `DELETE FROM docker_images WHERE created_at < NOW() - INTERVAL '1 day' * $1`

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

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

	if rowsAffected > 0 {
		fmt.Printf("Deleted %d old docker image records\n", rowsAffected)
	}

	return nil
}

// GetDockerStats returns statistics about Docker images across all agents
func (q *DockerQueries) GetDockerStats() (*models.DockerStats, error) {
	var stats models.DockerStats

	// Get total images
	err := q.db.QueryRow("SELECT COUNT(*) FROM docker_images").Scan(&stats.TotalImages)
	if err != nil {
		return nil, fmt.Errorf("failed to get total docker images: %w", err)
	}

	// Get images with updates
	err = q.db.QueryRow("SELECT COUNT(*) FROM docker_images WHERE current_version != available_version").Scan(&stats.UpdatesAvailable)
	if err != nil {
		return nil, fmt.Errorf("failed to get docker images with updates: %w", err)
	}

	// Get critical updates
	err = q.db.QueryRow("SELECT COUNT(*) FROM docker_images WHERE severity = 'critical' AND current_version != available_version").Scan(&stats.CriticalUpdates)
	if err != nil {
		return nil, fmt.Errorf("failed to get critical docker updates: %w", err)
	}

	// Get agents with Docker images
	err = q.db.QueryRow("SELECT COUNT(DISTINCT agent_id) FROM docker_images").Scan(&stats.AgentsWithContainers)
	if err != nil {
		return nil, fmt.Errorf("failed to get agents with docker images: %w", err)
	}

	return &stats, nil
}