package org.dromara.resource.junziqian.openai;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import java.io.BufferedReader;
import java.io.DataOutputStream;
import java.io.InputStreamReader;
import java.net.HttpURLConnection;
import java.net.URL;

@Service
public class ChatGPTUtils
{
    private static final Logger log;
    static String baseUrl;
    static String key;

    public static void main(final String[] args) throws Exception {
        final ChatGPTUtils chatGPTUtils = new ChatGPTUtils();
        final String q1 = "中国的个人所得税征收标准是什么";
        final ChatResponse response1 = chatGPTUtils.chat("123", "你是一个智能助手", "", "", q1);
        final String assistantPreContent = response1.getContent();
        final String q2 = "起征点是多少";
        final ChatResponse response2 = chatGPTUtils.chat("123", "你是一个智能助手", q1, assistantPreContent, q2);
        System.out.println(response2);
    }

    public ChatResponse chat(final String user, final String systemContent, final String userPreContent, final String assistantPreContent, final String userContent) {
        final ChatResponse response = new ChatResponse();
        final String modelId = "gpt-4-1106-preview";
        final String apiUrl = ChatGPTUtils.baseUrl + "/v1/chat/completions";
        final JSONObject data = new JSONObject();
        data.set("user", (Object)user);
        data.set("model", (Object)modelId);
        final ChatMessage chatMessage = new ChatMessage();
        final JSONArray messages = chatMessage.buildMessage(systemContent, userPreContent, assistantPreContent, "", userContent);
        data.set("messages", (Object)messages);
        data.set("max_tokens", (Object)4096);
        final JSONObject result = this.sendPost(apiUrl, data.toString());
        if (!result.containsKey((Object)"error")) {
            response.setContent(result.getStr("content"));
        }
        else {
            response.setError(result.getStr("error"));
        }
        return response;
    }

    public ChatResponse vision(final String user, final String systemContent, final String imageUrl, final String userContent) {
        final ChatResponse response = new ChatResponse();
        final String modelId = "gpt-4-vision-preview";
        final String apiUrl = ChatGPTUtils.baseUrl + "/v1/chat/completions";
        final JSONObject data = new JSONObject();
        data.set("user", (Object)user);
        data.set("model", (Object)modelId);
        final ChatMessage chatMessage = new ChatMessage();
        final JSONArray messages = chatMessage.buildMessage(systemContent, "", "", imageUrl, userContent);
        data.set("messages", (Object)messages);
        data.set("max_tokens", (Object)4096);
        final JSONObject result = this.sendPost(apiUrl, data.toString());
        if (!result.containsKey((Object)"error")) {
            response.setContent(result.getStr("content"));
        }
        else {
            response.setError(result.getStr("error"));
        }
        return response;
    }

    private JSONObject sendPost(final String apiUrl, final String data) {
        final JSONObject resultJSON = new JSONObject();
        String result = "";
        try {
            final URL url = new URL(apiUrl);
            final HttpURLConnection connection = (HttpURLConnection)url.openConnection();
            connection.setRequestMethod("POST");
            connection.setRequestProperty("Authorization", "Bearer " + ChatGPTUtils.key);
            connection.setRequestProperty("Content-Type", "application/json");
            connection.setRequestProperty("Accept", "application/json");
            connection.setDoOutput(true);
            try (final DataOutputStream outputStream = new DataOutputStream(connection.getOutputStream())) {
                ChatGPTUtils.log.info("postData:" + data);
                outputStream.write(data.getBytes());
                outputStream.flush();
            }
            final int responseCode = connection.getResponseCode();
            if (responseCode == 200) {
                try (final BufferedReader reader = new BufferedReader(new InputStreamReader(connection.getInputStream()))) {
                    final StringBuilder response = new StringBuilder();
                    String line;
                    while ((line = reader.readLine()) != null) {
                        response.append(line);
                    }
                    ChatGPTUtils.log.info("ChatGPT 响应:" + response.toString());
                    result = response.toString();
                }
            }
            else {
                ChatGPTUtils.log.info("ChatGPT 请求失败，响应码: " + responseCode);
            }
            connection.disconnect();
        }
        catch (Exception e) {
            e.printStackTrace();
            ChatGPTUtils.log.error("ChatGPT 网络请求错误", (Throwable)e);
            resultJSON.set("error", (Object)"ChatGPT 网络请求错误");
        }
        if (StrUtil.isNotBlank((CharSequence)result)) {
            final JSONObject jsonObject = JSONUtil.parseObj(result);
            final JSONArray choices = jsonObject.getJSONArray("choices");
            final JSONObject error = jsonObject.getJSONObject("error");
            if (choices != null && !jsonObject.containsKey((Object)"error")) {
                final JSONObject choice = choices.getJSONObject(0);
                final String role = choice.getJSONObject("message").getStr("role");
                final String content = choice.getJSONObject("message").getStr("content");
                resultJSON.set("role", (Object)role);
                resultJSON.set("content", (Object)content);
            }
            else if (error != null) {
                resultJSON.set("error", error.get((Object)"message"));
            }
        }
        return resultJSON;
    }

    static {
        log = LoggerFactory.getLogger((Class)ChatGPTUtils.class);
        ChatGPTUtils.baseUrl = "https://flag.smarttrot.com";
        ChatGPTUtils.key = "zk-b02953bb6f9045cf4c683a438350e8fb";
    }
}
