FloorVisualizer/internal/queue/redis.go

package queue

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

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

const (
	queueKey     = "jobs:queue"
	jobKeyPrefix = "jobs:"
	jobTTL       = 1 * time.Hour
)

// JobStatus represents the state of a generation job.
type JobStatus struct {
	JobID   string `json:"job_id"`
	Status  string `json:"status"`  // pending | processing | done | error
	Step    string `json:"step"`
	Message string `json:"message"`
	Result  string `json:"result,omitempty"`
	Error   string `json:"error,omitempty"`
}

// NewRedisClient creates a Redis client and pings to verify connectivity.
func NewRedisClient(addr, password string) (*redis.Client, error) {
	rdb := redis.NewClient(&redis.Options{
		Addr:     addr,
		Password: password,
		DB:       0,
	})
	if err := rdb.Ping(context.Background()).Err(); err != nil {
		return nil, fmt.Errorf("redis connect: %w", err)
	}
	return rdb, nil
}

// EnqueueJob pushes a job ID onto the queue and sets initial status.
func EnqueueJob(ctx context.Context, rdb *redis.Client, jobID string) error {
	pipe := rdb.Pipeline()
	pipe.RPush(ctx, queueKey, jobID)
	pipe.HSet(ctx, jobKey(jobID), "status", "pending")
	pipe.Expire(ctx, jobKey(jobID), jobTTL)
	_, err := pipe.Exec(ctx)
	return err
}

// DequeueJob blocks until a job is available, then pops and returns its ID.
func DequeueJob(ctx context.Context, rdb *redis.Client) (string, error) {
	result, err := rdb.BLPop(ctx, 0, queueKey).Result()
	if err != nil {
		return "", err
	}
	if len(result) < 2 {
		return "", fmt.Errorf("unexpected blpop result: %v", result)
	}
	return result[1], nil
}

// SetJobStatus updates the job status field.
func SetJobStatus(ctx context.Context, rdb *redis.Client, jobID, status string) error {
	return rdb.HSet(ctx, jobKey(jobID), "status", status).Err()
}

// SetJobProgress updates the step and message for a job.
func SetJobProgress(ctx context.Context, rdb *redis.Client, jobID, step, msg string) error {
	return rdb.HSet(ctx, jobKey(jobID), "step", step, "message", msg).Err()
}

// SetJobResult marks a job as done with the result image URL.
func SetJobResult(ctx context.Context, rdb *redis.Client, jobID, url string) error {
	return rdb.HSet(ctx, jobKey(jobID), "status", "done", "step", "done", "message", "Complete", "result", url).Err()
}

// SetJobError marks a job as errored with an error message.
func SetJobError(ctx context.Context, rdb *redis.Client, jobID, errMsg string) error {
	return rdb.HSet(ctx, jobKey(jobID), "status", "error", "step", "error", "message", errMsg, "error", errMsg).Err()
}

// GetJobStatus retrieves the full status of a job.
func GetJobStatus(ctx context.Context, rdb *redis.Client, jobID string) (*JobStatus, error) {
	m, err := rdb.HGetAll(ctx, jobKey(jobID)).Result()
	if err != nil {
		return nil, err
	}
	if len(m) == 0 {
		return nil, fmt.Errorf("job not found")
	}
	js := &JobStatus{
		JobID:   jobID,
		Status:  m["status"],
		Step:    m["step"],
		Message: m["message"],
		Result:  m["result"],
		Error:   m["error"],
	}
	return js, nil
}

// StoreJobPayload saves the job's input parameters (floor_id, pattern_code, room_code, input_path).
func StoreJobPayload(ctx context.Context, rdb *redis.Client, jobID string, payload map[string]string) error {
	data, _ := json.Marshal(payload)
	return rdb.HSet(ctx, jobKey(jobID), "payload", string(data)).Err()
}

// GetJobPayload retrieves the stored input parameters.
func GetJobPayload(ctx context.Context, rdb *redis.Client, jobID string) (map[string]string, error) {
	s, err := rdb.HGet(ctx, jobKey(jobID), "payload").Result()
	if err != nil {
		return nil, err
	}
	var p map[string]string
	json.Unmarshal([]byte(s), &p)
	return p, nil
}

func jobKey(id string) string {
	return jobKeyPrefix + id
}