package org.tzd.lm;

import java.io.*;
import java.net.HttpURLConnection;
import java.net.URI;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;

/**
 * 使用AI接口获取回复
 * @author tzdwindows 7
 */
public class LMApi {
    private static final String API_URL = "https://api.pearktrue.cn/api/deepseek/";
    private static final Gson GSON = new Gson();

    public static class Message {
        private String role;
        private String content;

        public Message(String role, String content) {
            this.role = role;
            this.content = content;
        }

        public String getRole() { return role; }
        public String getContent() { return content; }
        public void setRole(String role) { this.role = role; }
        public void setContent(String content) { this.content = content; }
    }

    /**
     * 调用AI接口获取回复
     * @param messages 消息列表
     * @return AI回复内容
     */
    public String getAIResponse(List<Message> messages) throws IOException {
        String jsonInput = GSON.toJson(Map.of("messages", messages));

        URL url = URI.create(API_URL).toURL();
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod("POST");
        conn.setRequestProperty("Content-Type", "application/json");
        conn.setDoOutput(true);

        try(OutputStream os = conn.getOutputStream()) {
            byte[] input = jsonInput.getBytes(StandardCharsets.UTF_8);
            os.write(input, 0, input.length);
        }

        int responseCode = conn.getResponseCode();
        if (responseCode == HttpURLConnection.HTTP_OK) {
            try (BufferedReader br = new BufferedReader(
                    new InputStreamReader(conn.getInputStream(), "GBK"))) {
                String response = readAllLines(br);
                Map<String, String> result = GSON.fromJson(response,
                        new TypeToken<Map<String, String>>(){}.getType());
                return result.get("message");
            }
        } else {
            throw new IOException("API请求失败，状态码：" + responseCode);
        }
    }

    private String readAllLines(BufferedReader reader) throws IOException {
        StringBuilder sb = new StringBuilder();
        String line;
        while ((line = reader.readLine()) != null) {
            sb.append(line);
        }
        return sb.toString();
    }

    public static void main(String[] args) {
        List<Message> messages = new ArrayList<>();
        messages.add(new Message("system", "You are a helpful assistant."));
        messages.add(new Message("user", "你好啊"));

        try {
            String reply = new LMApi().getAIResponse(messages);
            System.out.println(reply);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}