package com.backendsys.modules.sdk.comfyui.service.impl;

import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.StrUtil;
import com.backendsys.modules.common.config.redis.utils.RedisUtil;
import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskDao;
import com.backendsys.modules.sdk.comfyui.entity.ComfyuiTask;
import com.backendsys.modules.sdk.comfyui.rabbitmq.ComfyuiQueueConfig;
import com.backendsys.modules.sdk.comfyui.service.ComfyuiTaskService;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import org.springframework.amqp.rabbit.core.RabbitTemplate;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
public class ComfyuiTaskServiceImpl implements ComfyuiTaskService {

    @Value("${spring.application.name}")
    private String APPLICATION_NAME;

    @Autowired
    private RedisTemplate redisTemplate;
    @Autowired
    private RabbitTemplate rabbitTemplate;
    @Autowired
    private ComfyuiQueueConfig config;

    @Autowired
    private ComfyuiTaskDao comfyuiTaskDao;

    /**
     * 推送主队列任务 (以用户为队列单位)
     * 当用户提交新任务、或某任务执行结束后，都调用它
     */
    @Override
    @Transactional
    public void tryPushNext(long user_id) {
        tryPushNext(user_id, null);
    }
    @Override
    @Transactional
    public void tryPushNext(long user_id, String custom_params) {

        System.out.println("[推送主队列任务]");
        // 1. 如果该用户已在主队列里，直接返回
        String QUEUE_LOCK_KEY = APPLICATION_NAME + ":comfyui:queue:lock:user:" + Convert.toStr(user_id);
        System.out.println("-- (Redis) QUEUE_LOCK_KEY = " + QUEUE_LOCK_KEY);
        if (Boolean.TRUE.equals(redisTemplate.hasKey(QUEUE_LOCK_KEY))) {
            return;
        }

        // 2. 抢令牌（原子）
        Boolean ok = redisTemplate.opsForValue().setIfAbsent(QUEUE_LOCK_KEY, "1", Duration.ofMinutes(10));
        if (!Boolean.TRUE.equals(ok)) {
            return; // 抢失败（并发极小概率）
        }

        // 3. 该用户下一条未进主队列的任务
        LambdaQueryWrapper<ComfyuiTask> wrapperTask = new LambdaQueryWrapper<>();
        wrapperTask.eq(ComfyuiTask::getUser_id, user_id);
        wrapperTask.eq(ComfyuiTask::getIn_master, -1);
        List<ComfyuiTask> taskListByUser = comfyuiTaskDao.selectList(wrapperTask);
        if (taskListByUser == null || taskListByUser.isEmpty()) {
            System.out.println("用户没有任务了");
            redisTemplate.delete(String.format(QUEUE_LOCK_KEY, user_id));
            return;
        } else {
            System.out.println("用户队列数量: " + taskListByUser.size());
        }

        ComfyuiTask taskByUser = taskListByUser.get(0);
        wrapperTask.eq(ComfyuiTask::getId, taskByUser.getId());

        // 4. 打标 (进主队列)
        ComfyuiTask entity = new ComfyuiTask();
        entity.setIn_master(1);
        taskByUser.setIn_master(1);
        taskByUser.setCustom_params(custom_params);
        comfyuiTaskDao.update(entity, wrapperTask);

//        System.out.println("-- 在这里开始执行任务吗？ --");
//
//        String client_id = taskByUser.getClient_id();
//        System.out.println("task_id: " + taskByUser.getId());
//        System.out.println("client_id: " + client_id);
//        String prompt_object = taskByUser.getGenerate_request();
//        System.out.println("prompt_object: " + StrUtil.sub(prompt_object, 0, 200));
//
////        // [ComfyUI] 执行任务
////        JSONObject prompt_object = JSONUtil.parseObj(prompt);
////        Mono<ComfyuiResponse> cfPromptResponseMono = comfyuiService.prompt(client_id, prompt_object);
////        ComfyuiResponse response = cfPromptResponseMono.block();
////        response.setClient_id(client_id);

        // 5. 发送消息 (rabbitmq)
        System.out.println("-- 推送成功: (rabbitmq) task_id: " + taskByUser.getId());
        rabbitTemplate.convertAndSend(config.EXCHANGE, config.ROUTING_KEY, taskByUser);
        System.out.println("--------------------------------------------------------");

    }

}
