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

import cn.hutool.core.convert.Convert;
import cn.hutool.core.map.MapUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.backendsys.modules.common.config.redis.utils.RedisUtil;
import com.backendsys.modules.common.config.security.utils.SecurityUtil;
import com.backendsys.modules.crt.dao.CrtGenerateImageDao;
import com.backendsys.modules.crt.entity.CrtGenerateImage;
import com.backendsys.modules.sdk.comfyui.enums.TaskStatusEnums;
import com.backendsys.modules.sdk.comfyui.rabbitmq.ComfyuiQueueConfig;
import com.backendsys.modules.sdk.comfyui.rabbitmq.ComfyuiRabbitListener;
import com.backendsys.modules.sdk.comfyui.rabbitmq.DeliveryTagHolder;
import com.backendsys.modules.sdk.comfyui.rabbitmq.TagCtx;
import com.backendsys.modules.sdk.comfyui.service.ComfyuiSocketService;
import com.backendsys.modules.sdk.comfyui.service.ComfyuiTaskService;
import com.backendsys.modules.sdk.comfyui.utils.ComfyuiUtil;
import com.backendsys.modules.sse.entity.SseResponse;
import com.backendsys.modules.sse.entity.SseResponseEnum;
import com.backendsys.modules.sse.utils.SseUtil;
import com.backendsys.modules.upload.entity.SysFileResult;
import com.backendsys.modules.upload.service.SysFileService;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.GetResponse;
import org.springframework.amqp.rabbit.core.RabbitTemplate;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import org.springframework.web.reactive.socket.client.ReactorNettyWebSocketClient;
import org.springframework.web.reactive.socket.client.WebSocketClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.netty.http.client.HttpClient;

import java.net.URI;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.amqp.rabbit.connection.Connection;

@Service
public class ComfyuiSocketServiceImpl implements ComfyuiSocketService {

    @Autowired
    private SseUtil sseUtil;

    @Autowired
    private RedisUtil redisUtil;
    @Autowired
    private RabbitTemplate rabbitTemplate;


    @Autowired
    private ComfyuiQueueConfig config;
    @Autowired
    private ComfyuiUtil comfyuiUtil;
    @Autowired
    private ComfyuiTaskService comfyuiTaskService;
    @Autowired
    private SysFileService sysFileService;
    @Autowired
    private CrtGenerateImageDao crtGenerateImageDao;
    @Autowired
    private DeliveryTagHolder deliveryTagHolder;


    @Value("${spring.application.name}")
    private String APPLICATION_NAME;
    @Value("${comfyui.host}")
    private String COMFYUI_HOST;
    @Value("${comfyui.token}")
    private String COMFYUI_TOKEN;

    @Value("${comfyui.output}")
    private String COMFYUI_OUTPUT;

    /**
     * 管理多个连接
     */
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    /**
     * 创建带有认证头的 WebSocketClient
     * @param token 认证令牌
     * @return 配置好的 WebSocketClient
     */
    private WebSocketClient createWebSocketClientWithToken(String token) {
        HttpClient httpClient = HttpClient.create()
            .headers(headers -> headers.add("Authorization", "Bearer " + token))
            .responseTimeout(Duration.ofSeconds(30));  // 30秒超时
        return new ReactorNettyWebSocketClient(httpClient);
    }

    /**
     * [ComfyUI] 创建 WebSocket 监听连接
     */
    @Override
    public Mono<Void> connect(String client_id, Integer port) {

        String wsUrl =  "ws://" + COMFYUI_HOST + ":" + port + "/ws";
        return Mono.defer(() -> {
            if (sessions.containsKey(client_id)) {
                return Mono.error(new IllegalStateException("Connection already exists for client: " + client_id));
            }
            // 动态创建带有认证头的客户端
            WebSocketClient clientWithAuth = createWebSocketClientWithToken(COMFYUI_TOKEN);
            return clientWithAuth.execute(URI.create(wsUrl + "?clientId=" + client_id), session -> {
                // 保存会话
                sessions.put(client_id, session);
                // 接收消息
                Flux<String> incomingMessages = session.receive()
                    .map(WebSocketMessage::getPayloadAsText)
                    .doOnNext(message -> {
                        System.out.println("(doOnNext) Received from " + client_id + ": " + message);
                    })
                    .doOnError(e -> {
                        System.err.println("(doOnError) Error for " + client_id + ": " + e.getMessage());
                    })
                    .doFinally(signal -> {
                        System.out.println("(doFinally) Connection closed for " + client_id + ": " + signal);
                        sessions.remove(client_id);
                    });
                // 需要返回一个Mono<Void>来表示处理完成
                return incomingMessages.then();
            });
        });
    }

    /**
     * [ComfyUI] 创建 WebSocket 监听连接 (转发到 SSE)
     * - client_id: -
     * - port:      端口 (comfyui)
     * - is_save:   是否转存到 (cos/tos)
     * - params:    额外参数
     */
    @Override
    public Mono<Void> connectToSse(String client_id, Integer port) {
        return connectToSse(client_id, port, false, null);
    }
    @Override
    public Mono<Void> connectToSse(String client_id, Integer port, Boolean is_save) {
        return connectToSse(client_id, port, is_save, null);
    }
    @Override
    public Mono<Void> connectToSse(String client_id, Integer port, Boolean is_save, Map<String, Object> params) {
        return connectToSse(client_id, port, is_save, params, null);
    }
    // 如果是在 websocket 周期，则需要手动输入 user_id
    @Override
    public Mono<Void> connectToSse(String client_id, Integer port, Boolean is_save, Map<String, Object> params, Long input_user_id) {

        // 由于 Websocket 获取不到上下文信息，所以 user_id 不能在 socket 周期获取
        Long user_id = input_user_id != null ? input_user_id : SecurityUtil.getUserId();

        String wsUrl =  "ws://" + COMFYUI_HOST + ":" + port + "/ws";
        return Mono.defer(() -> {
            if (sessions.containsKey(client_id)) {
                return Mono.error(new IllegalStateException("Connection already exists for client: " + client_id));
            }
            // 动态创建带有认证头的客户端
            WebSocketClient clientWithAuth = createWebSocketClientWithToken(COMFYUI_TOKEN);
            String wsUrlFull = wsUrl + "?clientId=" + client_id + "&userId=" + user_id;
            System.out.println("Listen: " + wsUrlFull);
            return clientWithAuth.execute(URI.create(wsUrlFull), session -> {
                // 保存会话
                sessions.put(client_id, session);

                Flux<String> incomingMessages = session.receive()
                        .map(WebSocketMessage::getPayloadAsText)
                        .doOnNext(message -> {
                            System.out.println("(doOnNext) Received from " + client_id + ": " + message);

                            JSONObject data = JSONUtil.parseObj(message);
                            JSONObject dataChildren = JSONUtil.parseObj(data.get("data"));

                            Long drama_project_storyboard_id = Convert.toLong(params.get("drama_project_storyboard_id"));
                            dataChildren.put("drama_project_storyboard_id", drama_project_storyboard_id);
                            data.put("data", dataChildren);

                            // == [任务执行完成] =======================================================================
                            // { "type": "executed", .. }
                            String type = Convert.toStr(data.get("type"));
                            if (TaskStatusEnums.EXECUTED.getValue().equals(type)) {

                                // 在任务结束后(SSE)(ComfyuiSocketService.connectToSse):
                                // 1) 手动确认消息(ACK)
                                // 2) 清除队列锁 (QUEUE_LOCK_KEY)
                                // 3) 尝试把下一条任务推进主队列







                                // -- [RabbitMQ][完成当前队列任务，放行下一个队列] ----------------------------------------
                                try {
                                    System.out.println("-- 准备放行队列 --");
//                                    Connection connection = rabbitTemplate.getConnectionFactory().createConnection();
//                                    Channel channel = connection.createChannel(false);
//                                    GetResponse resp = channel.basicGet(config.QUEUE, false);

                                    // 已经被 Listener 的这里拿不到?
                                    // 这里是空的?
                                    TagCtx tag = deliveryTagHolder.take(client_id);
                                    Channel channel = tag.getChannel();

                                    // 1) 手动确认消息 (ACK)
                                    System.out.println("-- 手动确认消息 (ACK) ----------");
//                                    channel.basicAck(resp.getEnvelope().getDeliveryTag(), false);
                                    System.out.println("client_id = " + client_id + ", tag = " + tag);
                                    channel.basicAck(tag.getDeliveryTag(), false);
                                    channel.close();


                                    // 2) 清除队列锁 (QUEUE_LOCK_KEY)
                                    String QUEUE_LOCK_KEY = APPLICATION_NAME + ":comfyui:queue:lock:user:" + Convert.toStr(user_id);
                                    System.out.println("-- 清除队列锁 (QUEUE_LOCK_KEY): " + QUEUE_LOCK_KEY);
                                    redisUtil.delete(QUEUE_LOCK_KEY);



                                } catch (Exception e) {
                                    System.out.println(e.getMessage());
                                }
                                // ----------------------------------------------------------------------------------








                                JSONObject output = JSONUtil.parseObj(dataChildren.get("output"));

                                // -- [生成图片 -> 转存图片 -> 新增记录] ------------------------------------------------
                                // 由于图片地址不是公开的，需要加 Token 访问，因此不能公开返回原始图片地址，比如：
                                // http://43.128.1.201:8000/api/view?filename=fenjing_1_00012_.png&token=$2b$12$.MR4qGaFetN1FPQzbfyIrehsyjnPJ12xAZhR/l7KZpLkUPQTCG4gy
                                // 因此，直接在 ComfyUI 的物理目录上构建 nginx 静态资源访问目录
                                Object imagesObj = output.get("images");
                                if (imagesObj != null) {

                                    // 3) 尝试把下一条任务推进主队列 (生图) ----------------------------------------------
                                    System.out.println("-- 尝试把下一条任务推进主队列 (生图) --");
                                    Map<String, Object> map = new HashMap<>();
                                    map.put("drama_project_storyboard_id", drama_project_storyboard_id);
                                    map.put("user_id", user_id);
                                    String custom_params = JSONUtil.toJsonStr(map);

                                    System.out.println("custom_params = " + custom_params);

                                    comfyuiTaskService.tryPushNext(user_id, custom_params);
                                    // -----------------------------------------------------------------------------

                                    String prompt_id = Convert.toStr(dataChildren.get("prompt_id"));

                                    // 临时图片相对路径
                                    JSONArray images = JSONUtil.parseArray(imagesObj);
                                    // 返回完整的临时图片路径
                                    List<String> images_temp_path_full = imagesToFullpath(images, port);
                                    System.out.println("images_temp_path_full = " + images_temp_path_full);
                                    output.put("images_temp_path", images_temp_path_full);

                                    if (is_save) {
                                        // 创建一个 CompletableFuture 来执行异步任务
                                        CompletableFuture.runAsync(() -> {
                                            // 返回转存后的图片路径
                                            List<String> images_path = imagesToStoryboard(user_id, prompt_id, images_temp_path_full, drama_project_storyboard_id);
                                            output.put("images_path", images_path);
                                        });
                                    }

                                    dataChildren.put("output", output);
                                    data.put("data", dataChildren);

                                    // [DB] 更新任务状态
                                    CompletableFuture.runAsync(() -> {
                                        comfyuiUtil.updateComfyuiTask(prompt_id, JSONUtil.toJsonStr(dataChildren), 2);
                                    });

                                }

                            }
                            // ========================================================================

                            sseUtil.send(user_id, new SseResponse(SseResponseEnum.COMFYUI, data).toJsonStr());
                        })
                        .doOnError(e -> {
                            System.err.println("(doOnError) Error for " + client_id + ": " + e.getMessage());
                            sseUtil.send(user_id, new SseResponse(SseResponseEnum.COMFYUI, e.getMessage()).toJsonStr());
                        })
                        .doFinally(signal -> {
                            System.out.println("(doFinally) Connection closed for " + client_id + ": " + signal);
                            sseUtil.send(user_id, new SseResponse(SseResponseEnum.COMFYUI, signal).toJsonStr());
                            sessions.remove(client_id);
                            System.out.println("---------------------------------------------------------");
                        });
                // 需要返回一个Mono<Void>来表示处理完成
                return incomingMessages.then();
            });
        });
    }

    /**
     * [ComfyUI] 断开 WebSocket 监听连接
     * @param client_id 客户端ID
     * @return Mono<Void> 表示断开操作
     */
    @Override
    public Mono<Void> disconnect(String client_id) {
        return Mono.fromRunnable(() -> {
            WebSocketSession session = sessions.get(client_id);
            if (session != null) {
                System.out.println("disconnect success! client_id: " + client_id);

                Long user_id = SecurityUtil.getUserId();
                if (user_id != null) {
                    sseUtil.send(user_id, new SseResponse(SseResponseEnum.COMFYUI, "disconnect success! client_id: " + client_id).toJsonStr());
                }

                session.close().subscribe();
                sessions.remove(client_id);
            }
        });
    }

    /**
     * 获取所有活动的连接ID
     * @return 连接ID集合
     */
    public Flux<String> getActiveConnections() {
        return Flux.fromIterable(sessions.keySet());
    }

    // 返回临时图片完整路径
    private List<String> imagesToFullpath(JSONArray images, Integer port) {
        List<String> images_path = new ArrayList<>();
        if (images.size() > 0) {
            for (int i = 0; i < images.size(); i++) {
                JSONObject image = images.getJSONObject(i);
                String filename = image.getStr("filename");
//                String filepath = "http://o.daogu.ai/" + port + "/" + filename;
                String filepath = "http://" + COMFYUI_OUTPUT + "/" + port + "/" + filename;
                images_path.add(filepath);
            }
        }
        return images_path;
    }

    // 转存图片
    private List<String> imagesToStoryboard(Long user_id, String prompt_id, List<String> images_temp_path_full, Long drama_project_storyboard_id) {

        List<String> images_path = new ArrayList<>();
        if (images_temp_path_full.size() > 0) {

            images_temp_path_full.stream().forEach(temp_image -> {

                // -- [图片转存储存桶] -------------------------------------
                SysFileResult result = sysFileService.urlToUploadFile(temp_image, user_id);
                System.out.println("urlToUploadFile (result) = " + result);

                String image = result.getUrl();
                images_path.add(image);

                // -- [记录到生成图片记录表] --------------------------------
                if (drama_project_storyboard_id != null) {
                    CrtGenerateImage entity = new CrtGenerateImage();
                    entity.setUser_id(user_id);
                    entity.setDrama_project_storyboard_id(drama_project_storyboard_id);
                    entity.setPrompt_id(prompt_id);
                    entity.setUrl_origin(temp_image);
                    entity.setUrl(image);
                    entity.setTarget(result.getTarget());
                    entity.setObject_key(result.getKey());
                    crtGenerateImageDao.insert(entity);
                }
                // ------------------------------------------------------

            });

        }
        return images_path;
    }

}
