package com.backendsys.modules.sdk.comfyui.rabbitmq;

import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskDao;
import com.backendsys.modules.sdk.comfyui.entity.ComfyuiResponse;
import com.backendsys.modules.sdk.comfyui.entity.ComfyuiTask;
import com.backendsys.modules.sdk.comfyui.service.ComfyuiService;
import com.backendsys.modules.sdk.comfyui.service.ComfyuiSocketService;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.GetResponse;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.rabbit.annotation.RabbitListener;
import org.springframework.amqp.support.AmqpHeaders;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Lazy;
import org.springframework.messaging.handler.annotation.Header;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 自定义监听器
 */
@Component
@Lazy(false)
public class ComfyuiRabbitListener {

    @Autowired
    private ComfyuiQueueConfig config;
    @Autowired
    private DeliveryTagHolder deliveryTagHolder;

    @Autowired
    private ComfyuiService comfyuiService;
    @Autowired
    private ComfyuiTaskDao comfyuiTaskDao;


    // [监听] 手动 ACK
    // (从 ComfyuiTaskService.tryPushNext 过来)
    @RabbitListener(id = "ComfyuiContainer", queues = "comfyui.queue", ackMode = "MANUAL")
    public void receive(Message message, Channel channel, @Header(AmqpHeaders.DELIVERY_TAG) long deliveryTag) throws IOException {
        try {

            System.out.println("[Rabbitmq] 开始监听队列");

            // 2. 业务处理 (提交Comfyui任务)
            String body = new String(message.getBody(), StandardCharsets.UTF_8);
            // {"id": 24, "user_id": 1, "client_id": "..", "generate_request" :"..", ..}

            // 收到的参数:
            JSONObject body_object = JSONUtil.parseObj(body);
            String prompt = body_object.getStr("generate_request");
            String client_id = body_object.getStr("client_id");
            String custom_params = body_object.getStr("custom_params");

            JSONObject prompt_object = JSONUtil.parseObj(prompt);

            // [ComfyUI] 发起任务
            Mono<ComfyuiResponse> cfPromptResponseMono = comfyuiService.prompt(client_id, prompt_object, custom_params);
            ComfyuiResponse response = cfPromptResponseMono.block();
            System.out.println("[Comfyui][发起任务] response = " + response);

            // -- [DB] 更新 任务ID、任务状态、任务URL、任务URL端口 --------------------------
            ComfyuiTask entity = new ComfyuiTask();
            entity.setExecute_prompt_id(response.getPrompt_id());
            entity.setExecute_url(response.getUrl());
            entity.setExecute_url_port(response.getUrl_port());
            entity.setExecute_status(1);    // 1: 正在生成中

            LambdaQueryWrapper<ComfyuiTask> wrapperTask = new LambdaQueryWrapper<>();
            wrapperTask.eq(ComfyuiTask::getClient_id, client_id);
            comfyuiTaskDao.update(entity, wrapperTask);
            // ------------------------------------------------------------------------

            // 2. 暂存 deliveryTag
            System.out.println("put client_id = " + client_id);
            deliveryTagHolder.put(client_id, deliveryTag, channel);

            /*
                外部调用ACK (确认消息)：
                try {
                    Connection connection = rabbitTemplate.getConnectionFactory().createConnection();
                    Channel channel = connection.createChannel(false);
                    GetResponse resp = channel.basicGet("demo.queue", false);
                    if (resp != null) {
                        channel.basicAck(resp.getEnvelope().getDeliveryTag(), false);
                    }
                    channel.close();
                    connection.close();
                } catch (Exception e) {
                    System.err.println("(RabbitListener) 处理失败: " + e.getMessage());
                }
             */

            // 3. 手动确认（deliveryTag + 是否批量）
            // channel.basicAck(message.getMessageProperties().getDeliveryTag(), false);

        } catch (Exception e) {
            System.err.println("(RabbitListener) 处理失败: " + e.getMessage());
            // 拒绝并重新入队
            channel.basicNack(message.getMessageProperties().getDeliveryTag(), false, true);
        }
    }


    // [监听] 死信队列
    @RabbitListener(id = "ComfyuiDlxContainer", queues = "comfyui.dlx.queue", ackMode = "MANUAL")
    public void handleComfyuiDlx(Message message, Channel channel) throws IOException {

        String body = new String(message.getBody());
        System.out.println("[Comfyui][RabbitMQ-死信] 收到：" + StrUtil.sub(body, 0, 1000));

        // 业务：记录日志 / 重发 / 报警 / 人工补偿
        channel.basicAck(message.getMessageProperties().getDeliveryTag(), false);
    }


    // 4. 并发消费（一条队列多线程）
    /*
    @ComfyuiRabbitListener(queues = "demo.queue", concurrency = "5-10")
    public void receive(String msg) { ... }
     */



}
