FloorVisualizer/internal/queue/redis.go

package queue

import (
	"context"
	"encoding/json"
	"fmt"
	"strconv"
	"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"`
	ElapsedMs int64            `json:"elapsed_ms,omitempty"`
	TimingsMs map[string]int64 `json:"timings_ms,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", "queued_at", nowMillis(), "updated_at", nowMillis())
	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 {
	fields := []any{"status", status, "updated_at", nowMillis()}
	if status == "processing" {
		fields = append(fields, "started_at", nowMillis())
	}
	return rdb.HSet(ctx, jobKey(jobID), fields...).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, "updated_at", nowMillis()).Err()
}

// SetJobTiming stores a single step duration and the total elapsed time so far.
func SetJobTiming(ctx context.Context, rdb *redis.Client, jobID, step string, duration, elapsed time.Duration) error {
	return rdb.HSet(ctx, jobKey(jobID),
		step+"_ms", duration.Milliseconds(),
		"elapsed_ms", elapsed.Milliseconds(),
		"updated_at", nowMillis(),
	).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, "updated_at", nowMillis()).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, "updated_at", nowMillis()).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"],
		ElapsedMs: parseInt64(m["elapsed_ms"]),
		TimingsMs: map[string]int64{},
	}
	for _, step := range []string{"check", "mask", "inpaint", "cleanup", "total"} {
		if value := parseInt64(m[step+"_ms"]); value > 0 {
			js.TimingsMs[step] = value
		}
	}
	if len(js.TimingsMs) == 0 {
		js.TimingsMs = nil
	}
	return js, nil
}

// StoreJobPayload saves the job's input parameters and render options.
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
}

func nowMillis() int64 {
	return time.Now().UnixMilli()
}

func parseInt64(raw string) int64 {
	value, _ := strconv.ParseInt(raw, 10, 64)
	return value
}