package com.mortal.bigmodelr.test;

import com.mortal.bigmodelr.mapper.visualModelMapper;

import javax.annotation.Resource;
import java.io.*;
import java.net.*;
import java.nio.file.*;
import java.util.*;
import java.util.concurrent.*;

//批量缺陷图片下载
public class BatchImageDownloader {



    /**
     * 批量下载图片
     * @param imageUrls 图片URL列表
     * @param saveDir 本地保存目录
     * @param threadCount 线程数
     * @param timeout 超时时间(毫秒)
     */
    public static void batchDownloadImages(List<String> imageUrls, String saveDir,
                                           int threadCount, int timeout) {
        // 创建线程池
        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
        List<Future<DownloadResult>> futures = new ArrayList<>();

        try {
            // 创建保存目录
            Path dirPath = Paths.get(saveDir);
            if (!Files.exists(dirPath)) {
                Files.createDirectories(dirPath);
            }

            // 提交下载任务
            for (String url : imageUrls) {
                Future<DownloadResult> future = executor.submit(() -> {
                    try {
                        String fileName = extractFileName(url);
                        Path savePath = dirPath.resolve(fileName);
                        return downloadSingleImage(url, savePath.toString(), timeout);
                    } catch (Exception e) {
                        return new DownloadResult(url, false, e.getMessage());
                    }
                });
                futures.add(future);
            }

            // 统计结果
            int successCount = 0;
            for (Future<DownloadResult> future : futures) {
                DownloadResult result = future.get();
                if (result.success) {
                    successCount++;
                    System.out.printf("下载成功: %s -> %s%n", result.url, result.message);
                } else {
                    System.err.printf("下载失败: %s, 原因: %s%n", result.url, result.message);
                }
            }

            System.out.printf("批量下载完成! 成功: %d, 失败: %d%n",
                    successCount, imageUrls.size() - successCount);
        } catch (Exception e) {
            System.err.println("批量下载出错: " + e.getMessage());
        } finally {
            executor.shutdown();
        }
    }

    /**
     * 下载单个图片
     */
    private static DownloadResult downloadSingleImage(String imageUrl, String savePath, int timeout)
            throws IOException {
        HttpURLConnection conn = null;
        try {
            URL url = new URL(imageUrl);
            conn = (HttpURLConnection) url.openConnection();
            conn.setRequestMethod("GET");
            conn.setConnectTimeout(timeout);
            conn.setReadTimeout(timeout);

            // 检查响应码
            int responseCode = conn.getResponseCode();
            if (responseCode != HttpURLConnection.HTTP_OK) {
                return new DownloadResult(imageUrl, false,
                        "HTTP错误码: " + responseCode);
            }

            // 获取文件类型
            String contentType = conn.getContentType();
            if (!contentType.startsWith("image/")) {
                return new DownloadResult(imageUrl, false,
                        "不是图片类型: " + contentType);
            }

            // 下载文件
            try (InputStream in = conn.getInputStream()) {
                Files.copy(in, Paths.get(savePath), StandardCopyOption.REPLACE_EXISTING);
                return new DownloadResult(imageUrl, true, savePath);
            }
        } finally {
            if (conn != null) conn.disconnect();
        }
    }

    /**
     * 从URL提取文件名
     */
    private static String extractFileName(String url) {
        String fileName = url.substring(url.lastIndexOf('/') + 1);
        if (fileName.isEmpty()) {
            fileName = "image_" + System.currentTimeMillis() + ".jpg";
        }

        // 移除URL参数
        if (fileName.contains("?")) {
            fileName = fileName.substring(0, fileName.indexOf('?'));
        }

        return fileName;
    }

    /**
     * 下载结果封装类
     */
    private static class DownloadResult {
        String url;
        boolean success;
        String message;

        DownloadResult(String url, boolean success, String message) {
            this.url = url;
            this.success = success;
            this.message = message;
        }
    }

//    public static void main(String[] args) {
//        List<String> imageUrls = visualModelMapper.selectImgUrl();
//
//        batchDownloadImages(imageUrls, "C:\\Users\\王寻\\Desktop\\1111\\大模型\\模型图片集", DEFAULT_THREADS, DEFAULT_TIMEOUT);
//    }
}