001/*
002 *  Copyright (c) 2023-2026, Agents-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.agentsflex.core.model.client.impl;
017
018import com.agentsflex.core.model.chat.ChatConfig;
019import com.agentsflex.core.model.chat.log.ChatMessageLogger;
020import com.agentsflex.core.model.client.OkHttpClientUtil;
021import com.agentsflex.core.model.client.StreamClient;
022import com.agentsflex.core.model.client.StreamClientListener;
023import com.agentsflex.core.util.StringUtil;
024import okhttp3.*;
025import org.jetbrains.annotations.NotNull;
026import org.slf4j.Logger;
027import org.slf4j.LoggerFactory;
028
029import java.io.BufferedReader;
030import java.io.IOException;
031import java.io.InputStreamReader;
032import java.util.Map;
033
034public class DnjsonClient implements StreamClient, Callback {
035
036    private static final Logger log = LoggerFactory.getLogger(DnjsonClient.class);
037    private static final MediaType JSON_TYPE = MediaType.parse("application/json; charset=utf-8");
038
039    private OkHttpClient okHttpClient;
040    private StreamClientListener listener;
041    private ChatConfig config;
042    private boolean isStop = false;
043
044    public DnjsonClient() {
045        this(OkHttpClientUtil.buildDefaultClient());
046    }
047
048    public DnjsonClient(OkHttpClient okHttpClient) {
049        if (okHttpClient == null) {
050            throw new IllegalArgumentException("OkHttpClient must not be null");
051        }
052        this.okHttpClient = okHttpClient;
053    }
054
055    public OkHttpClient getOkHttpClient() {
056        return okHttpClient;
057    }
058
059    public void setOkHttpClient(OkHttpClient okHttpClient) {
060        this.okHttpClient = okHttpClient;
061    }
062
063    @Override
064    public void start(String url, Map<String, String> headers, String payload, StreamClientListener listener, ChatConfig config) {
065        if (isStop) {
066            throw new IllegalStateException("DnjsonClient has been stopped and cannot be reused.");
067        }
068
069        this.listener = listener;
070        this.config = config;
071        this.isStop = false;
072
073        Request.Builder builder = new Request.Builder().url(url);
074        if (headers != null && !headers.isEmpty()) {
075            headers.forEach(builder::addHeader);
076        }
077
078        RequestBody body = RequestBody.create(payload, JSON_TYPE);
079        Request request = builder.post(body).build();
080
081        ChatMessageLogger.logRequest(config, payload);
082
083        if (this.listener != null) {
084            try {
085                this.listener.onStart(this);
086            } catch (Exception e) {
087                log.warn("Error in listener.onStart", e);
088                return; // 可选：是否继续请求？
089            }
090        }
091
092        // 发起异步请求
093        okHttpClient.newCall(request).enqueue(this);
094    }
095
096    @Override
097    public void stop() {
098        // 注意：OkHttp 的 Call 无法取消已开始的 onResponse
099        // 所以 stop() 主要用于标记状态，防止后续回调处理
100        markAsStopped();
101    }
102
103
104    @Override
105    public void onFailure(@NotNull Call call, @NotNull IOException e) {
106        try {
107            if (listener != null && !isStop) {
108                Throwable error = Util.getFailureThrowable(e, null);
109                listener.onFailure(this, error);
110            }
111        } catch (Exception ex) {
112            log.warn("Error in listener.onFailure", ex);
113        } finally {
114            markAsStopped();
115        }
116    }
117
118    @Override
119    public void onResponse(@NotNull Call call, @NotNull Response response) throws IOException {
120        try {
121            if (!response.isSuccessful()) {
122                if (listener != null && !isStop) {
123                    Throwable error = Util.getFailureThrowable(null, response);
124                    listener.onFailure(this, error);
125                }
126                return;
127            }
128
129            ResponseBody body = response.body();
130            if (body == null || isStop) {
131                return;
132            }
133
134            // 使用 try-with-resources 确保流关闭
135            try (ResponseBody responseBody = body;
136                 BufferedReader reader = new BufferedReader(new InputStreamReader(responseBody.byteStream()))) {
137
138                String line;
139                while ((line = reader.readLine()) != null) {
140                    if (isStop) break; // 支持中途 stop()
141
142                    if (!StringUtil.hasText(line)) continue;
143
144                    String jsonLine = StringUtil.isJsonObject(line) ? line : "{" + line + "}";
145
146                    if (listener != null && !isStop) {
147                        try {
148                            ChatMessageLogger.logResponse(config, jsonLine);
149                            listener.onMessage(this, jsonLine);
150                        } catch (Exception e) {
151                            log.warn("Error in listener.onMessage", e);
152                        }
153                    }
154                }
155            }
156        } finally {
157            markAsStopped();
158        }
159    }
160
161
162    private void markAsStopped() {
163        if (isStop) return;
164        synchronized (this) {
165            if (isStop) return;
166            isStop = true;
167            if (listener != null) {
168                try {
169                    listener.onStop(this);
170                } catch (Exception e) {
171                    log.warn("Error in listener.onStop", e);
172                }
173            }
174        }
175    }
176}