package com.dora.service.impl;

import com.dora.service.QueueService;
import lombok.RequiredArgsConstructor;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Service;

import java.util.Objects;

@Service
@RequiredArgsConstructor
public class QueueServiceImpl implements QueueService {

    private final RedisTemplate<String, String> redisTemplate;

    private static final String QUEUE_KEY_PREFIX = "ai_task_queue:";
    private static final String PROCESSING_COUNT_KEY_PREFIX = "ai_task_processing_count:";

    @Override
    public void enqueue(String modelName, String taskNo) {
        redisTemplate.opsForList().rightPush(getQueueKey(modelName), taskNo);
    }

    @Override
    public String dequeue(String modelName) {
        return redisTemplate.opsForList().leftPop(getQueueKey(modelName));
    }

    @Override
    public long getQueueLength(String modelName) {
        Long size = redisTemplate.opsForList().size(getQueueKey(modelName));
        return size != null ? size : 0;
    }

    @Override
    public long incrementProcessingCount(String modelName) {
        return Objects.requireNonNull(redisTemplate.opsForValue().increment(getProcessingCountKey(modelName)));
    }

    @Override
    public long decrementProcessingCount(String modelName) {
        return Objects.requireNonNull(redisTemplate.opsForValue().decrement(getProcessingCountKey(modelName)));
    }

    @Override
    public int getProcessingCount(String modelName) {
        String countStr = redisTemplate.opsForValue().get(getProcessingCountKey(modelName));
        return countStr != null ? Integer.parseInt(countStr) : 0;
    }

    @Override
    public long removeFromQueue(String modelName, String taskNo) {
        // 使用 LREM 命令从列表中移除指定元素
        // count = 0 表示移除所有匹配的元素
        Long removed = redisTemplate.opsForList().remove(getQueueKey(modelName), 0, taskNo);
        return removed != null ? removed : 0;
    }

    private String getQueueKey(String modelName) {
        return QUEUE_KEY_PREFIX + modelName;
    }

    private String getProcessingCountKey(String modelName) {
        return PROCESSING_COUNT_KEY_PREFIX + modelName;
    }
}
