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

import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.backendsys.modules.common.Filter.WebClientFilter;
import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskDao;
import com.backendsys.modules.sdk.comfyui.entity.*;
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 org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import java.util.*;
import java.util.function.Consumer;

@Service
public class ComfyuiServiceImpl implements ComfyuiService {

    @Value("${comfyui.host}")
    private String COMFYUI_HOST;
    @Value("${comfyui.token}")
    private String COMFYUI_TOKEN;
    @Value("${comfyui.ports}")
    private String COMFYUI_PORTS_STR;
    @Value("${comfyui.is-save}")
    private Boolean IS_SAVE;

    @Autowired
    private ComfyuiTaskDao comfyuiTaskDao;
    @Autowired
    private ComfyuiSocketService comfyuiSocketService;

    private WebClient webClient;
    public WebClient getWebClient() {
        if (webClient == null) {
            webClient = WebClient.builder().filter(WebClientFilter.logFilter).build();
        }
        return webClient;
    }

    /**
     * 获取 Token 拼接好的 HttpHeaders
     */
    public Consumer<HttpHeaders> getHttpHeaders() {
        Consumer<HttpHeaders> httpHeaders = (headers) -> {
            headers.add("Authorization", "Bearer " + COMFYUI_TOKEN);
        };
        return httpHeaders;
    }


    /**
     * [ComfyUI] 查询任务队列
     */
    public Mono<ComfyuiQueue> getQueue(Integer port) {
        String url = "http://" + COMFYUI_HOST + ":" + port + "/queue";
        WebClient webClient = getWebClient();
        return webClient.get()
                .uri(url)
                .headers(getHttpHeaders())
                .accept(MediaType.APPLICATION_JSON)
                .exchangeToMono(response -> response.bodyToMono(ComfyuiQueue.class));
    }

    /**
     * 获取最佳可用端口
     */
    @Override
    public Integer getBestPort() {

//        List<Integer> PORTS = Arrays.asList(8000, 8001, 8002, 8003, 8004, 8005, 8006, 8007);
        List<Integer> PORTS = Convert.toList(Integer.class, StrUtil.split(COMFYUI_PORTS_STR, ","));

        List<ComfyuiQueueItem> queueCountsList = new ArrayList<>();

        // 查询所有端口的队列情况
        Boolean is_empty_queue = false;
        for (int i = 0; i < PORTS.size(); i++) {

            Integer port = PORTS.get(i);

            ComfyuiQueueItem queueCountsItem = new ComfyuiQueueItem();
            queueCountsItem.setPort(port);

            System.out.println("-- 请求端口: " + port + " ---------------");

            // [WebClient] 远程查询 ComfyuiAPI 队列情况
            Mono<ComfyuiQueue> cfQueueMono = getQueue(port);
            ComfyuiQueue response = cfQueueMono.block();
            if (response != null) {
                JSONArray queue_pending = response.getQueue_pending();
                JSONArray queue_running = response.getQueue_running();
                Integer count = queue_pending.size() + queue_running.size();
                System.out.println("port: " + port + ", queueCount: " + count);

                queueCountsItem.setCount(count);
                queueCountsList.add(0, queueCountsItem);  // 从前面添加记录

                if (queue_pending.size() == 0 && queue_running.size() == 0) {
                    System.out.println("发现空的队列，停止遍历: " + port);
                    is_empty_queue = true;
                    break;  // 中断循环
                }

            } else {
                System.out.println("port: " + port + ", response is null. ");
            }

        }

        // 没有空队列时
        // - 再次遍历，找出 count 最小值的行数据；
        // - 将最小值行数据排序到最前面；
        if (!is_empty_queue) {
            queueCountsList.sort(Comparator.comparing(ComfyuiQueueItem::getCount));
            System.out.println("-------------------------------");
            System.out.println("没有空队列，将队列数量最小值的行数据排序到最前面: ");
            System.out.println("port = " + queueCountsList.get(0).getPort() + ", count = " + queueCountsList.get(0).getCount());
        }

        System.out.println("-- 最终选择的端口: ----------------");
        System.out.println("port = " + queueCountsList.get(0).getPort() + ", count = " + queueCountsList.get(0).getCount());

        return queueCountsList.isEmpty() ? null : queueCountsList.get(0).getPort();
    }


    /**
     * [ComfyUI] 执行任务
     */
    @Override
    public Mono<ComfyuiResponse> prompt(String client_id, JSONObject prompt) {
        return prompt(client_id, prompt, null);
    }
    @Override
    public Mono<ComfyuiResponse> prompt(String client_id, JSONObject prompt, String custom_params) {

        // 获取最优端口
        Integer port = getBestPort();

        System.out.println("自定义参数 (custom_params) = " + custom_params);

        System.out.println("[ComfyUI] 开始监听 WecSocket 事件");
        JSONObject custom_params_object = (StrUtil.isNotEmpty(custom_params)) ? JSONUtil.parseObj(custom_params) : new JSONObject();
        // -- [ComfyUI] 创建 WebSocket 监听进度 --------------------------
        Map<String, Object> params = new LinkedHashMap<>();
        Long drama_project_storyboard_id = Convert.toLong(custom_params_object.get("drama_project_storyboard_id"));
        Long user_id = Convert.toLong(custom_params_object.get("user_id"));
        if (drama_project_storyboard_id != null) {
            params.put("drama_project_storyboard_id", drama_project_storyboard_id);
        }
        // - is_save: 是否转存到 cos/tos
        comfyuiSocketService.connectToSse(client_id, port, IS_SAVE, params, user_id).subscribe();
        // -------------------------------------------------------------
        // 在任务结束后(SSE)(ComfyuiSocketService.connectToSse):
        // 1) 手动确认消息(ACK)
        // 2) 清除队列锁 (QUEUE_LOCK_KEY)
        // 3) 尝试把下一条任务推进主队列

        ComfyuiRequest bodyValue = new ComfyuiRequest();
        bodyValue.setClient_id(client_id);
        bodyValue.setPrompt(prompt);

        String url = "http://" + COMFYUI_HOST + ":" + port + "/prompt";
        WebClient webClient = getWebClient();
        return webClient.post()
                .uri(url)
                .headers(getHttpHeaders())
                .accept(MediaType.APPLICATION_JSON)
                .bodyValue(bodyValue)
                .exchangeToMono(response -> response.bodyToMono(ComfyuiResponse.class))
                .map(r -> {
                    r.setUrl(url);
                    r.setUrl_port(port);
                    return r;
                })
                .onErrorResume(e -> {
                    // 捕获所有异常（包括上面抛出的 RuntimeException）
                    ComfyuiResponse response = new ComfyuiResponse();
                    response.setNode_errors(e.getMessage());

                    // -- [DB] 更新任务错误信息 --------------------------------------------------
                    ComfyuiTask entity = new ComfyuiTask();
                    entity.setReason(e.getMessage());
                    entity.setExecute_status(3);
                    LambdaQueryWrapper<ComfyuiTask> wrapperTask = new LambdaQueryWrapper<>();
                    wrapperTask.eq(ComfyuiTask::getClient_id, client_id);
                    comfyuiTaskDao.update(entity, wrapperTask);
                    // ------------------------------------------------------------------------

                    return Mono.just(response);
                });

    }

}
