heroagent/pkg/jobsmanager/redis.go

package jobsmanager

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/freeflowuniverse/heroagent/pkg/tools"
	"github.com/redis/go-redis/v9"
)

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

// NewRedisClient creates a new Redis client
func NewRedisClient(addr string, isUnixSocket bool) (*RedisClient, error) {
	// Determine network type
	networkType := "tcp"
	if isUnixSocket {
		networkType = "unix"
	}

	// 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 &RedisClient{
		client: client,
		ctx:    ctx,
	}, nil
}

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

// StoreJob stores a job in Redis
func (r *RedisClient) StoreJob(job *Job) error {
	// Convert job to JSON
	jobJSON, err := job.ToJSON()
	if err != nil {
		return err
	}

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

	return nil
}

// GetJob retrieves a job from Redis by ID
func (r *RedisClient) GetJob(jobID interface{}) (*Job, error) {
	// For interface{} type, we need more information to construct the storage key
	// Two approaches:
	// 1. For numeric JobID only (no circle or topic), use legacy format
	// 2. For full key with pattern "circleID:topic:jobID", parse out components

	var storageKey string

	// Handle different types of jobID
	switch id := jobID.(type) {
	case uint32:
		// Legacy format for backward compatibility
		storageKey = fmt.Sprintf("jobsmanager:%d", id)
	case string:
		// Check if this is a composite key (circleID:topic:jobID)
		parts := strings.Split(id, ":")
		if len(parts) == 3 {
			// This is a composite key
			circleID := tools.NameFix(parts[0])
			topic := tools.NameFix(parts[1])

			// Try to convert last part to uint32
			var numericID uint32
			if _, err := fmt.Sscanf(parts[2], "%d", &numericID); err == nil {
				storageKey = fmt.Sprintf("herojobs:%s:%s:%d", circleID, topic, numericID)
			} else {
				// Legacy string ID format in composite key
				storageKey = fmt.Sprintf("herojobs:%s:%s:%s", circleID, topic, parts[2])
			}
		} else {
			// Try to convert string to uint32 (legacy format)
			var numericID uint32
			if _, err := fmt.Sscanf(id, "%d", &numericID); err == nil {
				storageKey = fmt.Sprintf("jobsmanager:%d", numericID)
			} else {
				// Legacy string ID format
				storageKey = fmt.Sprintf("jobsmanager:%s", id)
			}
		}
	default:
		return nil, fmt.Errorf("unsupported job ID type: %T", jobID)
	}

	// Get job from Redis
	jobJSON, err := r.client.HGet(r.ctx, storageKey, "job").Result()
	if err != nil {
		if err == redis.Nil {
			return nil, fmt.Errorf("job not found: %v", jobID)
		}
		return nil, fmt.Errorf("failed to get job: %w", err)
	}

	// Parse job JSON
	job, err := NewJobFromJSON(jobJSON)
	if err != nil {
		return nil, err
	}

	return job, nil
}

// DeleteJob deletes a job from Redis
func (r *RedisClient) DeleteJob(jobID interface{}) error {
	// Handle different types of jobID similar to GetJob
	var storageKey string

	switch id := jobID.(type) {
	case uint32:
		// Legacy format for backward compatibility
		storageKey = fmt.Sprintf("jobsmanager:%d", id)
	case string:
		// Check if this is a composite key (circleID:topic:jobID)
		parts := strings.Split(id, ":")
		if len(parts) == 3 {
			// This is a composite key
			circleID := tools.NameFix(parts[0])
			topic := tools.NameFix(parts[1])

			// Try to convert last part to uint32
			var numericID uint32
			if _, err := fmt.Sscanf(parts[2], "%d", &numericID); err == nil {
				storageKey = fmt.Sprintf("herojobs:%s:%s:%d", circleID, topic, numericID)
			} else {
				// Legacy string ID format in composite key
				storageKey = fmt.Sprintf("herojobs:%s:%s:%s", circleID, topic, parts[2])
			}
		} else {
			// Try to convert string to uint32 (legacy format)
			var numericID uint32
			if _, err := fmt.Sscanf(id, "%d", &numericID); err == nil {
				storageKey = fmt.Sprintf("jobsmanager:%d", numericID)
			} else {
				// Legacy string ID format
				storageKey = fmt.Sprintf("jobsmanager:%s", id)
			}
		}
	default:
		return fmt.Errorf("unsupported job ID type: %T", jobID)
	}

	// Delete job from Redis
	err := r.client.Del(r.ctx, storageKey).Err()
	if err != nil {
		return fmt.Errorf("failed to delete job: %w", err)
	}

	return nil
}

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

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

	return nil
}

// QueueSize returns the size of a queue
func (r *RedisClient) QueueSize(circleID, topic string) (int64, error) {
	// Apply name fixing to CircleID and Topic
	fixedCircleID := tools.NameFix(circleID)
	fixedTopic := tools.NameFix(topic)
	queueKey := fmt.Sprintf("heroqueue:%s:%s", fixedCircleID, fixedTopic)

	// 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
}

// QueueEmpty empties a queue and deletes all corresponding jobs
func (r *RedisClient) QueueEmpty(circleID, topic string) error {
	// Apply name fixing to CircleID and Topic
	fixedCircleID := tools.NameFix(circleID)
	fixedTopic := tools.NameFix(topic)
	queueKey := fmt.Sprintf("heroqueue:%s:%s", fixedCircleID, fixedTopic)

	// Get all job IDs from queue
	jobIDs, err := r.client.LRange(r.ctx, queueKey, 0, -1).Result()
	if err != nil {
		return fmt.Errorf("failed to get job IDs from queue: %w", err)
	}

	// Delete all jobs
	for _, jobIDStr := range jobIDs {
		// Convert string ID to uint32 if possible
		var jobID uint32
		if _, err := fmt.Sscanf(jobIDStr, "%d", &jobID); err == nil {
			// New format with CircleID and Topic
			storageKey := fmt.Sprintf("herojobs:%s:%s:%d", fixedCircleID, fixedTopic, jobID)
			err := r.client.Del(r.ctx, storageKey).Err()
			if err != nil {
				return fmt.Errorf("failed to delete job %d: %w", jobID, err)
			}
		} else {
			// Handle legacy string IDs
			storageKey := fmt.Sprintf("jobsmanager:%s", jobIDStr)
			err := r.client.Del(r.ctx, storageKey).Err()
			if err != nil {
				return fmt.Errorf("failed to delete job %s: %w", jobIDStr, err)
			}
		}
	}

	// Empty queue
	err = r.client.Del(r.ctx, queueKey).Err()
	if err != nil {
		return fmt.Errorf("failed to empty queue: %w", err)
	}

	return nil
}

// QueueGet gets the last job from a queue without removing it
func (r *RedisClient) QueueGet(circleID, topic string) (*Job, error) {
	// Apply name fixing to CircleID and Topic
	fixedCircleID := tools.NameFix(circleID)
	fixedTopic := tools.NameFix(topic)
	queueKey := fmt.Sprintf("heroqueue:%s:%s", fixedCircleID, fixedTopic)

	// Get last job ID from queue
	jobID, err := r.client.LIndex(r.ctx, queueKey, -1).Result()
	if err != nil {
		if err == redis.Nil {
			return nil, fmt.Errorf("queue is empty")
		}
		return nil, fmt.Errorf("failed to get job ID from queue: %w", err)
	}

	// Get job with context about circle and topic
	compositeID := fmt.Sprintf("%s:%s:%s", circleID, topic, jobID)
	return r.GetJob(compositeID)
}

// QueueFetch gets and removes the last job from a queue
func (r *RedisClient) QueueFetch(circleID, topic string) (*Job, error) {
	// Apply name fixing to CircleID and Topic
	fixedCircleID := tools.NameFix(circleID)
	fixedTopic := tools.NameFix(topic)
	queueKey := fmt.Sprintf("heroqueue:%s:%s", fixedCircleID, fixedTopic)

	// Get and remove last job ID from queue
	jobID, err := r.client.RPop(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)
	}

	// Get job with context about circle and topic
	compositeID := fmt.Sprintf("%s:%s:%s", circleID, topic, jobID)
	return r.GetJob(compositeID)
}

// ListJobs lists job IDs by circle and topic
func (r *RedisClient) ListJobs(circleID, topic string) ([]uint32, error) {
	var pattern string

	// Apply name fixing to CircleID and Topic
	fixedCircleID := tools.NameFix(circleID)
	fixedTopic := tools.NameFix(topic)

	if circleID != "" && topic != "" {
		pattern = fmt.Sprintf("heroqueue:%s:%s", fixedCircleID, fixedTopic)
	} else if circleID != "" {
		pattern = fmt.Sprintf("heroqueue:%s:*", fixedCircleID)
	} else if topic != "" {
		pattern = fmt.Sprintf("heroqueue:*:%s", fixedTopic)
	} else {
		pattern = "heroqueue:*:*"
	}

	// Get all matching queue keys
	queueKeys, err := r.client.Keys(r.ctx, pattern).Result()
	if err != nil {
		return nil, fmt.Errorf("failed to list queues: %w", err)
	}

	var jobIDs []uint32
	for _, queueKey := range queueKeys {
		// Get all job IDs from queue
		stringIDs, err := r.client.LRange(r.ctx, queueKey, 0, -1).Result()
		if err != nil {
			return nil, fmt.Errorf("failed to get job IDs from queue %s: %w", queueKey, err)
		}

		// Convert string IDs to uint32
		for _, idStr := range stringIDs {
			var id uint32
			if _, err := fmt.Sscanf(idStr, "%d", &id); err == nil {
				jobIDs = append(jobIDs, id)
			} else {
				// Log but continue - this handles legacy string IDs that can't be converted
				fmt.Printf("Warning: Found job ID that couldn't be converted to uint32: %s\n", idStr)
			}
		}
	}

	return jobIDs, nil
}

// UpdateJobStatus updates the status of a job
func (r *RedisClient) UpdateJobStatus(jobID uint32, status JobStatus) error {
	// Get job
	job, err := r.GetJob(jobID)
	if err != nil {
		return err
	}

	// Update status
	job.Status = status

	// Update timestamps based on status
	now := time.Now().Unix()
	if status == JobStatusActive && job.TimeStart == 0 {
		job.TimeStart = now
	} else if (status == JobStatusDone || status == JobStatusError) && job.TimeEnd == 0 {
		job.TimeEnd = now
	}

	// Store updated job
	return r.StoreJob(job)
}

// UpdateJobResult updates the result of a job
func (r *RedisClient) UpdateJobResult(jobID uint32, result string) error {
	// Get job
	job, err := r.GetJob(jobID)
	if err != nil {
		return err
	}

	// Update result
	job.Result = result
	job.Status = JobStatusDone
	job.TimeEnd = time.Now().Unix()

	// Store updated job
	return r.StoreJob(job)
}

// UpdateJobError updates the error of a job
func (r *RedisClient) UpdateJobError(jobID uint32, errorMsg string) error {
	// Get job
	job, err := r.GetJob(jobID)
	if err != nil {
		return err
	}

	// Update error
	job.Error = errorMsg
	job.Status = JobStatusError
	job.TimeEnd = time.Now().Unix()

	// Store updated job
	return r.StoreJob(job)
}