package cn.com.poc.dify;

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.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.chat.ChatMessage;
import cn.com.gsst.dify_client.model.chat.ChatMessageResponse;
import cn.com.yict.framemax.core.spring.SingleContextInitializer;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
import org.springframework.test.context.web.WebAppConfiguration;

import java.io.IOException;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.CountDownLatch;

/**
 * @author alex.yao
 * @date 2025/9/15
 */
@RunWith(SpringJUnit4ClassRunner.class)
@ContextConfiguration(initializers = SingleContextInitializer.class)
@WebAppConfiguration
public class DifyTest {

    @Test
    public void test_chat() throws DifyApiException, IOException {
        // 创建聊天客户端
        DifyChatClient chatClient = DifyClientFactory.createChatClient("http://192.168.21.174:6060/v1", "app-pv9FlrZCMfYogqzmbSxwg16s");

        // 创建聊天消息
        ChatMessage message = ChatMessage.builder()
                .query("你好，请介绍一下自己")
                .user("user-123")
                .inputs(new LinkedHashMap<>())
                .responseMode(ResponseMode.BLOCKING)
                .build();

        // 发送消息并获取响应
        ChatMessageResponse response = chatClient.sendChatMessage(message);
        System.out.println("回复: " + response.getAnswer());
        System.out.println("会话ID: " + response.getConversationId());
        System.out.println("消息ID: " + response.getMessageId());
    }

    @Test
    public void test_chatStream() throws DifyApiException, IOException, InterruptedException {
        // 创建聊天客户端
        DifyChatClient chatClient = DifyClientFactory.createChatClient("https://dify.gsstcloud.com/v1", "app-YmVPoXmTh7ejERFbip3ADCYN");

        Map<String,Object> inputs = new LinkedHashMap<>();
        inputs.put("demand","Java实习");
        inputs.put("type","chat");
        inputs.put("jobInfomations","chat");
        inputs.put("fileinfos","{\"baseInfo\":{\"department\":\"技术部\",\"educationalRe\":\"本科及以上\",\"jobNature\":\"全职\",\"jobTitle\":\"AI应用工程师\",\"numberOfRecruits\":5,\"salaryRange\":\"15k-20k\",\"workExperience\":\"3-5年\",\"workLocation\":\"北京\"}}");
        // 创建聊天消息
        ChatMessage message = ChatMessage.builder()
                .query("帮我优化 职业画像")
                .user("user-123")
                .inputs(inputs)
                .responseMode(ResponseMode.STREAMING)
                .build();
        CountDownLatch countDownLatch = new CountDownLatch(1);

        // 发送消息并获取响应
        chatClient.sendChatMessageStream(message, new ChatStreamCallback() {
            @Override
            public void onMessage(MessageEvent event) {
                ChatStreamCallback.super.onMessage(event);
                System.out.println("answer:" + event.getAnswer());
            }

            @Override
            public void onMessageEnd(MessageEndEvent event) {
                ChatStreamCallback.super.onMessageEnd(event);
                countDownLatch.countDown();
                System.out.println("answer:" + "[DONE]");
            }
        });
        countDownLatch.await();
    }
}
