package com.vetti.common.ai.gpt;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.vetti.common.ai.gpt.service.OpenAiStreamListenerService;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * AI ChatGPT 对话
 */
@Component
public class OpenAiStreamClient {

    @Value("${chatGpt.apiKey}")
    private String apiKey;
    // API端点URL
    @Value("${chatGpt.apiUrl}")
    private String apiUrl;

    @Value("${chatGpt.model}")
    private String model;

    @Value("${chatGpt.modelQuestion}")
    private String modelQuestion;

    @Value("${chatGpt.role}")
    private String role;

    // 定义作为分割点的标点符号集合
    private final String punctuationStr = "。？！；.?!;";

    /**
     * 发送流式请求
     *
     * @param promptJson   提示词json数据集合
     * @param listener 流式响应监听器
     */
    public void streamChat(String promptJson, OpenAiStreamListenerService listener) {
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(60, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
        // 文本缓冲区
        StringBuffer bufferStr = new StringBuffer();

        // 构建请求参数
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", modelQuestion);
        requestBody.put("stream", true);
        // 构建消息
        if(StrUtil.isNotEmpty(promptJson)) {
            List<Map> promptList = JSONUtil.toList(promptJson, Map.class);
            Object[] objects = new Object[promptList.size()];
            for (int i = 0; i < promptList.size(); i++) {
                objects[i] = promptList.get(i);
            }
            //获取到的提示
            requestBody.put("messages", objects);
        }
        // 创建请求
        Request request = new Request.Builder()
                .url(apiUrl)
                .addHeader("Content-Type", "application/json")
                .addHeader("Authorization", "Bearer " + apiKey)
                .post(RequestBody.create(JSONUtil.toJsonStr(requestBody),
                        MediaType.parse("application/json; charset=utf-8")
                ))
                .build();
        // 发送异步请求
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                listener.onError(e);
            }

            @Override
            public void onResponse(Call call, Response response) throws IOException {
                if (!response.isSuccessful()) {
                    listener.onError(new IOException("Unexpected code " + response));
                    return;
                }
                // 处理流式响应
                ResponseBody body = response.body();
                if (body == null) {
                    listener.onError(new IOException("Response body is null"));
                    return;
                }
                try (var source = body.source()) {
                    // 逐行读取响应
                    while (!source.exhausted()) {
                        String line = source.readUtf8Line();
                        if (line == null || line.isEmpty()) continue;
                        // 处理SSE格式（去除"data: "前缀）
                        if (line.startsWith("data: ")) {
                            String data = line.substring(6).trim();
                            // 检查是否为结束标记
                            if (data.equals("[DONE]")) {
                                listener.onMessage(bufferStr.toString());
                                bufferStr.setLength(0);
                                listener.onComplete();
                                break;
                            }
                            // 解析JSON获取内容
                            try {
                                JSONObject json = JSONUtil.parseObj(data);
//                                System.out.println(json);
                                String content = json.getJSONArray("choices")
                                        .getJSONObject(0)
                                        .getJSONObject("delta")
                                        .getStr("content");

                                if (content != null && !content.isEmpty()) {
                                    //加入缓冲区
                                    bufferStr.append(content);
                                    if(punctuationStr.contains(content.trim())){
                                        //说明有标点啦,直接返回
                                        listener.onMessage(bufferStr.toString());
                                        bufferStr.setLength(0);
                                    }
                                }
//                                listener.onMessage(content);
                            } catch (Exception e) {
                                listener.onError(new IOException("Parse error: " + e.getMessage()));
                            }
                        }
                    }
                } finally {
                    response.close();
                }
            }
        });
    }

}
