package cc.mrbird.febs.mall.controller;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionResult;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;

public class HSController {
    // 请确保您已将 AK SK 分别存储在环境变量 VOLC_ACCESSKEY 和 VOLC_SECRETKEY中
    // 初始化Ark客户端，从环境变量中读取您的AKSK
    // 从环境变量中获取您的Key鉴权。此为默认方式，您可根据需要进行修改
    private static String ak = "AKLTZTQxZjMyZTUxMWJmNDEyNDkzNWExOGQ3ODllNzhhNmQ";
    private static String sk = "TmpFeE1qZ3haREExTW1JeE5HRTBZVGc1WlRRNVlqWXpORGd5TWpsak5HWQ==";
    private static String ep_id = "ep-20250724113103-fn6pg";
    // 此为默认路径，您可根据业务所在地域进行配置
    static String baseUrl = "https://ark.cn-beijing.volces.com/api/v3";
    static ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
    static Dispatcher dispatcher = new Dispatcher();
    static ArkService service = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).baseUrl(baseUrl).ak(ak).sk(sk).build();

    public static void main(String[] args) {
        System.out.println("\n----- standard request -----");
        // 2. 构建对话历史（AI提问→用户回答）
        List<ChatMessage> messages = Arrays.asList(
                // System消息：定义AI角色和行为
                ChatMessage.builder()
                        .role(ChatMessageRole.SYSTEM)
                        .content("你是同事相处顾问，任务流程：\n1) 先提出一个关于团队协作的问题\n2) 用户回答后，分析回答的不足\n3) 提供具体的改进建议和示例")
                        .build(),
                // Assistant消息：AI的初始问题
                ChatMessage.builder()
                        .role(ChatMessageRole.ASSISTANT)
                        .content("组员成为了销冠，组长的你该如何表达自己的称赞，同时给其他组员加油？")
                        .build(),
                // User消息：用户的回答
                ChatMessage.builder()
                        .role(ChatMessageRole.USER)
                        .content("干的漂亮。")
                        .build()
        );
//        List<ChatMessage> messages = Arrays.asList(
//                // System消息：定义AI角色和行为
//                ChatMessage.builder()
//                        .role(ChatMessageRole.SYSTEM)
//                        .content("\"content\": \"###\n" +
//                                "假如你是一位专业的彩票数据分析师，你将根据双色球历史开奖记录相关信息，来解决回答关于双色球历史开奖结果相关问题的任务。根据以下规则一步步执行：\n" +
//                                "1. 回答问题时需提供准确的红球号码。\n" +
//                                "2. 回答问题时需提供准确的蓝球号码。\n" +
//                                "3. 回答问题时需提供准确的开奖日期。\n" +
//                                "\n" +
//                                "参考例子：\n" +
//                                "无\n" +
//                                "\n" +
//                                "请回答问题：\n" +
//                                "关于双色球历史开奖结果的问题\n" +
//                                "输出：\n" +
//                                "\n" +
//                                "要求：\n" +
//                                "1 输出包含准确的红球号码、蓝球号码和开奖日期。\n" +
//                                "2 表述清晰准确。\n" +
//                                "###")
//                        .build(),
//                // Assistant消息：AI的初始问题
////                ChatMessage.builder()
////                        .role(ChatMessageRole.ASSISTANT)
////                        .content("组员成为了销冠，组长的你该如何表达自己的称赞，同时给其他组员加油？")
////                        .build(),
//                // User消息：用户的回答
//                ChatMessage.builder()
//                        .role(ChatMessageRole.USER)
//                        .content("第2025001期双色球开奖结果是什么？")
//                        .build()
//        );
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                // 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
                .model(ep_id)
                .messages(messages)
                .build();
        // 4. 发送请求并处理响应
//        System.out.println("\n----- standard request -----");
//        service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> System.out.println(choice.getMessage().getContent()));

        System.out.println("\n----- streaming request -----");
        service.streamChatCompletion(chatCompletionRequest)
                .doOnError(Throwable::printStackTrace)
                .blockingForEach(
                        choice -> {
                            if (choice.getChoices().size() > 0) {
                                System.out.print(choice.getChoices().get(0).getMessage().getContent());
                            }
                        }
                );

        service.shutdownExecutor();
    }

}