package cn.com.poc.common.service.impl;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import cn.com.gsst.dify_client.DifyChatClient;
import cn.com.gsst.dify_client.DifyClientFactory;
import cn.com.gsst.dify_client.callback.ChatStreamCallback;
import cn.com.gsst.dify_client.enums.FileTransferMethod;
import cn.com.gsst.dify_client.enums.FileType;
import cn.com.gsst.dify_client.enums.ResponseMode;
import cn.com.gsst.dify_client.event.MessageEndEvent;
import cn.com.gsst.dify_client.event.MessageEvent;
import cn.com.gsst.dify_client.exception.DifyApiException;
import cn.com.gsst.dify_client.model.DifyConfig;
import cn.com.gsst.dify_client.model.chat.ChatMessage;
import cn.com.gsst.dify_client.model.chat.ChatMessageResponse;
import cn.com.gsst.dify_client.model.file.FileInfo;
import cn.com.poc.common.config.DifyConfigProperties;
import cn.com.poc.common.service.DifyChatService;
import cn.com.poc.common.utils.JsonUtils;
import cn.com.poc.common.utils.SSEUtil;
import cn.com.poc.common.utils.StringUtils;
import cn.com.poc.thirdparty.resource.demand.ai.entity.largemodel.LargeModelDemandResult;
import cn.com.yict.framemax.core.exception.BusinessException;
import org.apache.commons.collections4.CollectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

/**
 * 通用Dify聊天服务，封装与Dify API交互的逻辑
 */
@Service
public class DifyChatServiceImpl implements DifyChatService {

    private final Logger logger = LoggerFactory.getLogger(DifyChatServiceImpl.class);

    @Autowired
    private DifyConfigProperties difyConfigProperties;

    /**
     * 阻塞式聊天请求
     */
    public ChatMessageResponse blockingChat(List<String> fileUrls,
                                            String question,
                                            Map<String, Object> data,
                                            String user,
                                            String apiKey) {

        // 从配置中获取API密钥，如果传入的apiKey为null或空，则从配置中查找
        String actualApiKey = apiKey;
        if (StringUtils.isEmpty(actualApiKey)) {
            actualApiKey = difyConfigProperties.getApiKeys().get(user);
        }

        DifyConfig difyConfig = DifyConfig.builder()
                .apiKey(actualApiKey)
                .readTimeout(200 * 1000)
                .baseUrl(difyConfigProperties.getBaseUrl())
                .build();
        DifyChatClient chatClient = DifyClientFactory.createChatClient(difyConfig);

        List<FileInfo> files = new ArrayList<>();
        if (CollectionUtils.isNotEmpty(fileUrls)) {
            for (String fileUrl : fileUrls) {
                files.add(FileInfo.builder()
                        .type(FileType.DOCUMENT)
                        .url(fileUrl)
                        .transferMethod(FileTransferMethod.REMOTE_URL)
                        .build());
            }
        }

        // 创建聊天消息
        ChatMessage message = ChatMessage.builder()
                .query(question)
                .user(user)
                .files(files)
                .inputs(data)
                .responseMode(ResponseMode.BLOCKING)
                .build();

        // 发送消息并获取响应
        try {
            ChatMessageResponse chatMessageResponse = chatClient.sendChatMessage(message);
            return chatMessageResponse;
        } catch (IOException | DifyApiException e) {
            logger.error("执行{}生成异常: {}", user, e.getMessage(), e);
            throw new BusinessException("执行" + user + "生成异常");
        }
    }

    /**
     * 流式聊天请求
     */
    public void streamChat(List<String> fileUrls,
                           String question,
                           Map<String, Object> data,
                           String user,
                           String apiKey) throws IOException, DifyApiException, InterruptedException {

        // 从配置中获取API密钥，如果传入的apiKey为null或空，则从配置中查找
        String actualApiKey = apiKey;
        if (StringUtils.isEmpty(actualApiKey)) {
            actualApiKey = difyConfigProperties.getApiKeys().get(user);
        }

        // 创建聊天客户端
        DifyChatClient chatClient =
                DifyClientFactory.createChatWorkflowClient(difyConfigProperties.getBaseUrl(), actualApiKey);

        List<FileInfo> files = new ArrayList<>();
        if (CollectionUtils.isNotEmpty(fileUrls)) {
            for (String fileUrl : fileUrls) {
                files.add(FileInfo.builder()
                        .type(FileType.DOCUMENT)
                        .url(fileUrl)
                        .transferMethod(FileTransferMethod.REMOTE_URL)
                        .build());
            }
        }

        // 创建聊天消息
        ChatMessage message = ChatMessage.builder()
                .query(question)
                .user(user)
                .files(files)
                .inputs(data)
                .responseMode(ResponseMode.STREAMING)
                .build();

        CountDownLatch countDownLatch = new CountDownLatch(1);
        ServletRequestAttributes servletRequestAttributes =
                (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletResponse response = servletRequestAttributes.getResponse();
        // 发送消息并获取响应
        chatClient.sendChatMessageStream(message, new ChatStreamCallback() {
            SSEUtil sseUtil = new SSEUtil(response);

            @Override
            public void onMessage(MessageEvent event) {
                try {
                    if (StringUtils.isNotEmpty(event.getAnswer())) {
                        LargeModelDemandResult largeModelDemandResult = new LargeModelDemandResult();
                        largeModelDemandResult.setCode("0");
                        largeModelDemandResult.setMessage(event.getAnswer());
                        sseUtil.send(JsonUtils.serialize(largeModelDemandResult));
                    }
                } catch (Exception e) {
                    logger.error("发送SSE消息时出错: {}", e.getMessage(), e);
                    sseUtil.completeByError(e.getMessage());
                    throw new RuntimeException(e);
                }
            }

            @Override
            public void onMessageEnd(MessageEndEvent event) {
                try {
                    sseUtil.send("[DONE]");
                } catch (IOException e) {
                    logger.error("发送SSE结束消息时出错: {}", e.getMessage(), e);
                    sseUtil.completeByError(e.getMessage());
                    throw new RuntimeException(e);
                } finally {
                    sseUtil.complete();
                }
                countDownLatch.countDown();
            }
        });
        countDownLatch.await(120, TimeUnit.SECONDS);
    }
}