/**
 * @Datetime : 2024-08-17 11:04:20
 * @Author : WANGKANG
 * @Email : 1686617586@qq.com
 * @Blog :
 * @File : VideoCapture.java
 * @brief : 视频转图片工具类
 */

package com.taais.biz.utils;

import java.io.IOException;
import java.io.File;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;

import org.bytedeco.javacv.Java2DFrameConverter;


import org.bytedeco.opencv.opencv_core.Stream;

public class VideoCapture {
    /**
     * 截取视频帧并保存为图片
     * 设置唯一视频名称，保存到指定的路径下
     *
     * @param filePath   视频文件路径
     * @param savePath__ 图片保存路径
     * @param fps        帧率
     * @throws IOException
     * @throws InterruptedException
     */
    public static void startCaputreFromFile(String filePath, String savePath__, Long fps)
        throws IOException {
        FFmpegFrameGrabber grabber = FFmpegFrameGrabber.createDefault(filePath);

        Path savePath = Paths.get(savePath__);
        if (Files.exists(savePath)) {
            System.out.println("目录已存在: " + savePath);
            System.out.println("清空里面的内容");
            try {
                deleteDirectory(savePath);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        Files.createDirectories(savePath);

        System.out.println("--- 开始截图 ---");
        grabber.start();
        long current_time_stamp = 0;
        long current_time_stamp_step = 1000000 / fps;
        int idx = 0;
        while (true) {
            Frame frame = grabber.grabImage();
            if (frame != null) {
                // System.out.println(frame.timestamp);
                if (frame.timestamp >= current_time_stamp) {
                    idx += 1;
                    current_time_stamp += current_time_stamp_step;
                    saveImage(frame, savePath__, idx);
                }
            } else {
                break;
            }
        }
        grabber.stop();
        grabber.release();
        System.out.println("--- 截图完成---");
    }

    public static void startCaputreFromUrl(String videoUrl, String savePath__, Long fps)
        throws IOException, InterruptedException {
        FFmpegFrameGrabber grabber = FFmpegFrameGrabber.createDefault(videoUrl);
        grabber.setOption("rtsp_transport", "tcp"); // 使用tcp的方式，不然会丢包很严重

        Path savePath = Paths.get(savePath__);
        if (Files.exists(savePath)) {
            System.out.println("目录已存在: " + savePath);
            System.out.println("清空里面的内容");
            try {
                deleteDirectory(savePath);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        Files.createDirectories(savePath);

        System.out.println("--- 开始截图 ---");
        grabber.start();
        long current_time_stamp = 0;
        long current_time_stamp_step = 1000000 / fps;
        int idx = 0;

        boolean isInterrupted = false;
        while (true) {
            // 判断是否被中断
            if (Thread.currentThread().isInterrupted()) {
                System.out.println("############### 线程中断1 #################");
                isInterrupted = true;
                break;
            }
            Frame frame = grabber.grabImage();
            if (frame != null) {
                // System.out.println(frame.timestamp);
                if (frame.timestamp >= current_time_stamp) {
                    idx += 1;
                    current_time_stamp += current_time_stamp_step;
                    saveImage(frame, savePath__, idx);
                }
            } else {
                break;
            }
        }
        grabber.stop();
        grabber.release();

        if (isInterrupted) {
            throw new InterruptedException("线程中断");
        }
    }

    public static boolean saveImage(Frame frame, String savePath, int idx) {
        try {
            Path outPath = Paths.get(savePath, String.format("%05d", idx) + /*"_" + System.currentTimeMillis() + */ ".jpg");
            // System.out.println("图片已保存：" + outPath);
            File outPut = new File(outPath.toString());
            ImageIO.write(frameToBufferedImage(frame), "jpg", outPut);
            return true;
        } catch (IOException e) {
            e.printStackTrace();
            return false;
        }
    }

    /**
     * 创建BufferedImage对象
     */
    public static BufferedImage frameToBufferedImage(Frame frame) {
        Java2DFrameConverter converter = new Java2DFrameConverter();
        BufferedImage bufferedImage = converter.getBufferedImage(frame);
        converter.close();
        return bufferedImage;
    }

    public static int getFramgeByTime(FFmpegFrameGrabber grabber, Long timestamp) throws Exception {
        grabber.setTimestamp(timestamp, true);
        Frame frame = grabber.grabFrame(false, true, true, false); // 抓取对应的帧
        if (frame != null) {
            return grabber.getFrameNumber();
        } else {
            throw new Exception("无法获取对应时间戳的帧");
        }
    }

    public static long getTimestampByFrame(FFmpegFrameGrabber grabber, int frame) throws Exception {
        grabber.setFrameNumber(frame);
        return grabber.getTimestamp();
    }

    // 递归删除目录及其内容
    private static void deleteDirectory(Path directory) throws Exception {
        if (Files.isDirectory(directory)) {
            // 遍历目录中的所有文件和子目录
            try (java.util.stream.Stream<Path> stream = Files.list(directory)) {
                stream.forEach(path -> {
                    try {
                        deleteDirectory(path);
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                });
            }
        }
        // 删除文件或空目录
        Files.delete(directory);
    }
}
