heroagent_go_old/pkg/servers/heroagent/redis.go

package heroagent

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

	"github.com/redis/go-redis/v9"
)

// RedisJobManager handles Redis operations for jobs
type RedisJobManager struct {
	client *redis.Client
	ctx    context.Context
}

// NewRedisJobManager creates a new Redis job manager
func NewRedisJobManager(tcpPort int, unixSocketPath string) (*RedisJobManager, error) {
	// Determine network type and address
	var networkType, addr string
	if unixSocketPath != "" {
		networkType = "unix"
		addr = unixSocketPath
	} else {
		networkType = "tcp"
		addr = fmt.Sprintf("localhost:%d", tcpPort)
	}

	// Create Redis client
	client := redis.NewClient(&redis.Options{
		Network:      networkType,
		Addr:         addr,
		DB:           0,
		DialTimeout:  5 * time.Second,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 5 * time.Second,
	})

	// Test connection
	ctx := context.Background()
	_, err := client.Ping(ctx).Result()
	if err != nil {
		return nil, fmt.Errorf("failed to connect to Redis: %w", err)
	}

	return &RedisJobManager{
		client: client,
		ctx:    ctx,
	}, nil
}

// Close closes the Redis client
func (r *RedisJobManager) Close() error {
	return r.client.Close()
}

// QueueKey returns the Redis queue key for a topic
func QueueKey(topic string) string {
	return fmt.Sprintf("heroqueue:%s", topic)
}

// StorageKey returns the Redis storage key for a job
func StorageKey(jobID uint32, topic string) string {
	return fmt.Sprintf("herojobs:%s:%d", topic, jobID)
}

// StoreJob stores a job in Redis
func (r *RedisJobManager) StoreJob(job *Job) error {
	// Convert job to JSON
	jobJSON, err := json.Marshal(job)
	if err != nil {
		return fmt.Errorf("failed to marshal job: %w", err)
	}

	// Store job in Redis
	storageKey := StorageKey(job.JobID, job.Topic)
	err = r.client.Set(r.ctx, storageKey, jobJSON, 0).Err()
	if err != nil {
		return fmt.Errorf("failed to store job in Redis: %w", err)
	}

	return nil
}

// EnqueueJob adds a job to its queue
func (r *RedisJobManager) EnqueueJob(job *Job) error {
	// Store the job first
	if err := r.StoreJob(job); err != nil {
		return err
	}

	// Add job ID to queue
	queueKey := QueueKey(job.Topic)
	err := r.client.RPush(r.ctx, queueKey, job.JobID).Err()
	if err != nil {
		return fmt.Errorf("failed to enqueue job: %w", err)
	}

	log.Printf("Job %d enqueued in Redis queue %s", job.JobID, queueKey)
	return nil
}

// GetJob retrieves a job from Redis
func (r *RedisJobManager) GetJob(jobID uint32, topic string) (*Job, error) {
	// Get job from Redis
	storageKey := StorageKey(jobID, topic)
	jobJSON, err := r.client.Get(r.ctx, storageKey).Result()
	if err != nil {
		if err == redis.Nil {
			return nil, fmt.Errorf("job not found: %d", jobID)
		}
		return nil, fmt.Errorf("failed to get job from Redis: %w", err)
	}

	// Parse job JSON
	job := &Job{}
	if err := json.Unmarshal([]byte(jobJSON), job); err != nil {
		return nil, fmt.Errorf("failed to unmarshal job: %w", err)
	}

	return job, nil
}

// DeleteJob deletes a job from Redis
func (r *RedisJobManager) DeleteJob(jobID uint32, topic string) error {
	// Delete job from Redis
	storageKey := StorageKey(jobID, topic)
	err := r.client.Del(r.ctx, storageKey).Err()
	if err != nil {
		return fmt.Errorf("failed to delete job from Redis: %w", err)
	}

	log.Printf("Job %d deleted from Redis", jobID)
	return nil
}

// FetchNextJob fetches the next job from a queue
func (r *RedisJobManager) FetchNextJob(topic string) (*Job, error) {
	queueKey := QueueKey(topic)

	// Get and remove first job ID from queue
	jobIDStr, err := r.client.LPop(r.ctx, queueKey).Result()
	if err != nil {
		if err == redis.Nil {
			return nil, fmt.Errorf("queue is empty")
		}
		return nil, fmt.Errorf("failed to fetch job ID from queue: %w", err)
	}

	// Convert job ID to uint32
	jobID, err := strconv.ParseUint(jobIDStr, 10, 32)
	if err != nil {
		return nil, fmt.Errorf("invalid job ID: %s", jobIDStr)
	}

	// Get job from Redis
	return r.GetJob(uint32(jobID), topic)
}

// ListQueues lists all job queues
func (r *RedisJobManager) ListQueues() ([]string, error) {
	// Get all queue keys
	queueKeys, err := r.client.Keys(r.ctx, "heroqueue:*").Result()
	if err != nil {
		return nil, fmt.Errorf("failed to list queues: %w", err)
	}

	// Extract topic names from queue keys
	topics := make([]string, 0, len(queueKeys))
	for _, queueKey := range queueKeys {
		// Extract topic from queue key (format: heroqueue:<topic>)
		topic := queueKey[10:] // Skip "heroqueue:"
		topics = append(topics, topic)
	}

	return topics, nil
}

// QueueSize returns the size of a queue
func (r *RedisJobManager) QueueSize(topic string) (int64, error) {
	queueKey := QueueKey(topic)

	// Get queue size
	size, err := r.client.LLen(r.ctx, queueKey).Result()
	if err != nil {
		return 0, fmt.Errorf("failed to get queue size: %w", err)
	}

	return size, nil
}

// UpdateJobStatus updates the status of a job in Redis
func (r *RedisJobManager) UpdateJobStatus(job *Job) error {
	// Update job in Redis
	return r.StoreJob(job)
}