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.file2text.source;
017
018
019import com.agentsflex.core.file2text.util.IOUtils;
020
021import java.io.*;
022import java.net.HttpURLConnection;
023import java.net.URL;
024import java.nio.file.Paths;
025import java.util.concurrent.atomic.AtomicBoolean;
026import java.util.regex.Matcher;
027import java.util.regex.Pattern;
028
029/**
030 * 从 HTTP/HTTPS URL 读取文档的输入源，支持缓存（避免重复请求）
031 * 自动判断使用内存缓存还是临时文件缓存
032 */
033public class HttpDocumentSource implements DocumentSource {
034    private static final int DEFAULT_CONNECT_TIMEOUT = 20_000;
035    private static final int DEFAULT_READ_TIMEOUT = 60_000;
036    private static final long MEMORY_THRESHOLD = 10 * 1024 * 1024; // 10MB 以内走内存
037
038    private final String url;
039    private final String providedFileName;
040    private final String mimeType;
041    private final int connectTimeout;
042    private final int readTimeout;
043    private final java.util.function.Consumer<HttpURLConnection> connectionCustomizer;
044
045    private volatile byte[] cachedBytes = null;
046    private volatile File tempFile = null;
047    private volatile String resolvedFileName = null;
048    private volatile String resolvedMimeType = null;
049    private final AtomicBoolean downloaded = new AtomicBoolean(false);
050
051    public HttpDocumentSource(String url) {
052        this(url, null, null, DEFAULT_CONNECT_TIMEOUT, DEFAULT_READ_TIMEOUT, null);
053    }
054
055    public HttpDocumentSource(String url, String fileName) {
056        this(url, fileName, null, DEFAULT_CONNECT_TIMEOUT, DEFAULT_READ_TIMEOUT, null);
057    }
058
059    public HttpDocumentSource(String url, String fileName, String mimeType) {
060        this(url, fileName, mimeType, DEFAULT_CONNECT_TIMEOUT, DEFAULT_READ_TIMEOUT, null);
061    }
062
063    public HttpDocumentSource(
064        String url,
065        String fileName,
066        String mimeType,
067        int connectTimeout,
068        int readTimeout,
069        java.util.function.Consumer<HttpURLConnection> connectionCustomizer
070    ) {
071        this.url = validateUrl(url);
072        this.providedFileName = fileName;
073        this.mimeType = mimeType;
074        this.connectTimeout = connectTimeout;
075        this.readTimeout = readTimeout;
076        this.connectionCustomizer = connectionCustomizer;
077    }
078
079    private String validateUrl(String url) {
080        try {
081            new URL(url).toURI();
082            return url;
083        } catch (Exception e) {
084            throw new RuntimeException("Invalid URL: " + url);
085        }
086    }
087
088    @Override
089    public String getFileName() {
090        if (resolvedFileName != null) {
091            return resolvedFileName;
092        }
093        synchronized (this) {
094            if (resolvedFileName == null) {
095                resolvedFileName = detectFileName();
096            }
097        }
098        return resolvedFileName;
099    }
100
101    private String detectFileName() {
102        // 1. 用户提供
103        if (providedFileName != null && !providedFileName.trim().isEmpty()) {
104            return sanitizeFileName(providedFileName);
105        }
106
107        // 2. 从 URL 路径提取
108        String fromUrl = extractFileNameFromUrl();
109        if (fromUrl != null) return fromUrl;
110
111        // 3. 从 Content-Disposition 提取（需要连接）
112        try {
113            HttpURLConnection conn = createConnection();
114            conn.setRequestMethod("HEAD"); // 只获取头
115            conn.connect();
116            String fromHeader = extractFileNameFromHeader(conn);
117            conn.disconnect();
118            if (fromHeader != null) return fromHeader;
119        } catch (IOException e) {
120            // 忽略
121        }
122
123        return "downloaded-file";
124    }
125
126    private String extractFileNameFromUrl() {
127        try {
128            URL urlObj = new URL(this.url);
129            String path = urlObj.getPath();
130            if (path != null && path.length() > 1) {
131                String name = Paths.get(path).getFileName().toString();
132                if (name.contains(".")) {
133                    return sanitizeFileName(name);
134                }
135            }
136        } catch (Exception e) {
137            // 忽略
138        }
139        return null;
140    }
141
142    private String extractFileNameFromHeader(HttpURLConnection conn) {
143        try {
144            String header = conn.getHeaderField("Content-Disposition");
145            if (header != null) {
146                Pattern pattern = Pattern.compile("filename\\s*=\\s*\"?([^\";]+)\"?");
147                Matcher matcher = pattern.matcher(header);
148                if (matcher.find()) {
149                    return sanitizeFileName(matcher.group(1));
150                }
151            }
152        } catch (Exception e) {
153            // 忽略
154        }
155        return null;
156    }
157
158    public static String sanitizeFileName(String filename) {
159        if (filename == null) return "unknown";
160        return filename
161            .replaceAll("[\\\\/:*?\"<>|]", "_")
162            .replaceAll("\\.\\.", "_")
163            .replaceAll("^\\s+|\\s+$", "")
164            .isEmpty() ? "file" : filename;
165    }
166
167    @Override
168    public String getMimeType() {
169        if (resolvedMimeType != null) {
170            return resolvedMimeType;
171        }
172        synchronized (this) {
173            if (resolvedMimeType == null) {
174                try {
175                    HttpURLConnection conn = createConnection();
176                    conn.setRequestMethod("HEAD");
177                    conn.connect();
178                    resolvedMimeType = conn.getContentType();
179                    conn.disconnect();
180                } catch (IOException e) {
181                    resolvedMimeType = mimeType; // fallback
182                }
183                if (resolvedMimeType == null) {
184                    resolvedMimeType = mimeType;
185                }
186            }
187        }
188        return resolvedMimeType;
189    }
190
191    @Override
192    public InputStream openStream() throws IOException {
193        downloadIfNeeded();
194        if (cachedBytes != null) {
195            return new ByteArrayInputStream(cachedBytes);
196        } else if (tempFile != null) {
197            return new FileInputStream(tempFile);
198        } else {
199            throw new IOException("No content available");
200        }
201    }
202
203    /**
204     * 下载一次，缓存结果
205     */
206    private void downloadIfNeeded() throws IOException {
207        if (downloaded.get()) return;
208
209        synchronized (this) {
210            if (downloaded.get()) return;
211
212            HttpURLConnection conn = createConnection();
213            conn.connect();
214
215            try {
216                int code = conn.getResponseCode();
217                if (code >= 400) {
218                    throw new IOException("HTTP " + code + " from " + url);
219                }
220
221                // 判断是否走内存 or 临时文件
222                long contentLength = conn.getContentLengthLong();
223                boolean useMemory = contentLength > 0 && contentLength <= MEMORY_THRESHOLD;
224
225                if (useMemory) {
226                    // 内存缓存
227                    this.cachedBytes = IOUtils.toByteArray(conn.getInputStream(), MEMORY_THRESHOLD);
228                } else {
229                    // 临时文件缓存
230                    this.tempFile = File.createTempFile("http-", ".cache");
231                    this.tempFile.deleteOnExit();
232                    try (FileOutputStream fos = new FileOutputStream(tempFile)) {
233                        IOUtils.copyStream(conn.getInputStream(), fos, Long.MAX_VALUE);
234                    }
235                }
236
237                // 更新 MIME（如果未指定）
238                if (this.resolvedMimeType == null) {
239                    this.resolvedMimeType = conn.getContentType();
240                }
241
242            } finally {
243                conn.disconnect();
244            }
245
246            downloaded.set(true);
247        }
248    }
249
250    private HttpURLConnection createConnection() throws IOException {
251        URL urlObj = new URL(this.url);
252        HttpURLConnection conn = (HttpURLConnection) urlObj.openConnection();
253        conn.setConnectTimeout(connectTimeout);
254        conn.setReadTimeout(readTimeout);
255        conn.setInstanceFollowRedirects(true);
256        conn.setRequestMethod("GET");
257        conn.setRequestProperty("User-Agent", "DocumentParser/1.0");
258        if (connectionCustomizer != null) {
259            connectionCustomizer.accept(conn);
260        }
261        return conn;
262    }
263
264    /**
265     * 获取缓存大小（用于调试）
266     */
267    public long getCachedSize() {
268        if (cachedBytes != null) return cachedBytes.length;
269        if (tempFile != null) return tempFile.length();
270        return 0;
271    }
272
273    /**
274     * 清理临时文件
275     */
276    public void cleanup() {
277        if (tempFile != null && tempFile.exists()) {
278            tempFile.delete();
279        }
280    }
281}