FFmpegConvert/src/main/java/net/knarcraft/ffmpegconverter/utility/FFMpegHelper.java

package net.knarcraft.ffmpegconverter.utility;

import net.knarcraft.ffmpegconverter.container.FFMpegCommand;
import net.knarcraft.ffmpegconverter.container.ProcessResult;
import net.knarcraft.ffmpegconverter.container.StreamProbeResult;
import net.knarcraft.ffmpegconverter.property.StreamType;
import net.knarcraft.ffmpegconverter.streams.AudioStream;
import net.knarcraft.ffmpegconverter.streams.OtherStream;
import net.knarcraft.ffmpegconverter.streams.StreamObject;
import net.knarcraft.ffmpegconverter.streams.StreamTag;
import net.knarcraft.ffmpegconverter.streams.SubtitleStream;
import net.knarcraft.ffmpegconverter.streams.VideoStream;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * A class which helps with ffmpeg probing and converting
 */
public final class FFMpegHelper {

    private static final String PROBE_SPLIT_CHARACTER = "øæåÆØå";

    private FFMpegHelper() {

    }

    /**
     * Gets streams from a file
     *
     * @param ffprobePath     <p>The path/command to ffprobe</p>
     * @param file            <p>The file to probe</p>
     * @param subtitleFormats <p>The extensions to accept for external subtitles</p>
     * @param audioFormats    <p>The extensions to accept for external audio files</p>
     * @return <p>A list of StreamObjects</p>
     * @throws IOException <p>If the process can't be readProcess</p>
     */
    @NotNull
    public static StreamProbeResult probeFile(@NotNull String ffprobePath, @NotNull File file,
                                              @NotNull List<String> subtitleFormats,
                                              @NotNull List<String> audioFormats) throws IOException {
        return parseStreams(ffprobePath, probeForStreams(ffprobePath, file), file, subtitleFormats, audioFormats);
    }

    /**
     * Creates a list containing all required arguments for converting a video to a web playable video
     *
     * @param executable <p>The executable to use (ffmpeg/ffprobe)</p>
     * @param files      <p>The files to execute on</p>
     * @return <p>A FFMPEG command for web-playable video</p>
     */
    @NotNull
    public static FFMpegCommand getFFMpegWebVideoCommand(@NotNull String executable, @NotNull List<File> files) {
        FFMpegCommand command = getFFMpegGeneralFileCommand(executable, files);
        command.addOutputFileOption("-vcodec", "h264");
        command.addOutputFileOption("-pix_fmt", "yuv420p");
        command.addOutputFileOption("-ar", "48000");
        command.addOutputFileOption("-movflags", "+faststart");
        command.addOutputFileOption("-map_metadata", "0");
        command.addOutputFileOption("-movflags", "+use_metadata_tags");
        return command;
    }

    /**
     * Creates a list containing command line arguments for a general file
     *
     * @param executable <p>The executable to use (ffmpeg/ffprobe)</p>
     * @param files      <p>The files to execute on</p>
     * @return <p>A basic FFMPEG command</p>
     */
    @NotNull
    public static FFMpegCommand getFFMpegGeneralFileCommand(@NotNull String executable, @NotNull List<File> files) {
        FFMpegCommand command = new FFMpegCommand(executable);
        command.addGlobalOption("-nostdin");
        for (File file : files) {
            command.addInputFile(file.getName());
        }
        command.addOutputFileOption("-map_metadata", "0");
        command.addOutputFileOption("-movflags", "+use_metadata_tags");
        return command;
    }

    /**
     * Starts and prints output of a process
     *
     * @param processBuilder <p>The process to run</p>
     * @param folder         <p>The folder the process should run in</p>
     * @param spacer         <p>The character(s) to use between each new line read</p>
     * @param write          <p>Whether to write the output directly instead of storing it</p>
     * @return <p>The result of running the process</p>
     * @throws IOException <p>If the process can't be readProcess</p>
     */
    @NotNull
    public static ProcessResult runProcess(@NotNull ProcessBuilder processBuilder, @Nullable File folder,
                                           @NotNull String spacer, boolean write) throws IOException {
        //Give the user information about what's about to happen
        OutputUtil.print("Command to be run: ");
        OutputUtil.println(processBuilder.command().toString());

        //Set directory and error stream
        if (folder != null) {
            processBuilder.directory(folder);
        }
        processBuilder.redirectErrorStream(true);

        Process process = processBuilder.start();
        BufferedReader processReader = new BufferedReader(new InputStreamReader(process.getInputStream()));
        StringBuilder output = new StringBuilder();
        while (process.isAlive()) {
            String read = readProcess(processReader, spacer);
            if (read.isEmpty()) {
                continue;
            }

            if (write) {
                OutputUtil.println(read);
            } else {
                OutputUtil.printDebug(read);
                output.append(read);
            }
        }
        try {
            int exitCode = process.waitFor();
            OutputUtil.println("Process finished with exit code: " + exitCode);
            return new ProcessResult(exitCode, output.toString());
        } catch (InterruptedException e) {
            return new ProcessResult(1, output.toString());
        }
    }

    /**
     * Adds arguments for converting a file to h264 using hardware acceleration
     *
     * @param command <p>The command to add the arguments to</p>
     * @param quality <p>The quality to encode. 0 = best, 51 = worst.</p>
     */
    public static void addH264HardwareEncoding(@NotNull FFMpegCommand command, int quality) {
        command.addOutputFileOption("-codec:v", "h264_nvenc");
        command.addOutputFileOption("-profile", "high");
        command.addOutputFileOption("-preset", "p7");
        command.addOutputFileOption("-crf", String.valueOf(quality));
    }

    /**
     * Adds arguments for converting a file to h265 using hardware acceleration
     *
     * @param command <p>The command to add the arguments to</p>
     * @param quality <p>The quality to encode. 0 = best, 51 = worst.</p>
     */
    public static void addH265HardwareEncoding(@NotNull FFMpegCommand command, int quality) {
        command.addOutputFileOption("-codec:v", "hevc_nvenc");
        command.addOutputFileOption("-profile", "main10");
        command.addOutputFileOption("-preset", "p7");
        command.addOutputFileOption("-tag:v", "hvc1");
        command.addOutputFileOption("-crf", String.valueOf(quality));
    }

    /**
     * Maps all streams in the given list to the output in the given command
     *
     * @param command <p>The command to add the mappings to</p>
     * @param streams <p>The streams to map</p>
     * @param <K>     <p>The type of stream object to map</p>
     */
    public static <K extends StreamObject> void mapAllStreams(@NotNull FFMpegCommand command, @NotNull List<K> streams) {
        for (StreamObject stream : streams) {
            mapStream(command, stream);
        }
    }

    /**
     * Maps the given stream to the given FFMPEG command's output
     *
     * @param command <p>The command to map the stream to</p>
     * @param stream  <p>The stream to map</p>
     */
    public static void mapStream(@NotNull FFMpegCommand command, @NotNull StreamObject stream) {
        command.addOutputFileOption("-map", String.format("%d:%d",
                stream.getInputIndex(), stream.getAbsoluteIndex()));
    }

    /**
     * Escapes special characters which can cause trouble for ffmpeg
     *
     * @param fileName <p>The filename to escape.</p>
     * @return <p>A filename with known special characters escaped.</p>
     */
    @NotNull
    public static String escapeSpecialCharactersInFileName(@NotNull String fileName) {
        return fileName.replaceAll("\\\\", "\\\\\\\\\\\\\\\\")
                .replaceAll("'", "'\\\\\\\\\\\\\''")
                .replaceAll("%", "\\\\\\\\\\\\%")
                .replaceAll(":", "\\\\\\\\\\\\:")
                .replace("]", "\\]")
                .replace("[", "\\[");
    }

    /**
     * Gets the nth stream from a list of streams
     *
     * @param streams <p>A list of streams</p>
     * @param n       <p>The index of the audio stream to get</p>
     * @return <p>The first audio stream found, or null if no audio streams were found</p>
     */
    public static <G extends StreamObject> G getNthSteam(@NotNull List<G> streams, int n) {
        if (n < 0) {
            throw new IllegalArgumentException("N cannot be negative!");
        }
        G stream = null;
        if (streams.size() > n) {
            stream = streams.get(n);
        } else if (!streams.isEmpty()) {
            stream = streams.get(0);
        }
        return stream;
    }

    /**
     * Gets a list of all streams in a file
     *
     * @param ffprobePath <p>The path/command to ffprobe.</p>
     * @param file        <p>The file to probe.</p>
     * @return <p>A list of streams.</p>
     * @throws IOException <p>If something goes wrong while probing.</p>
     */
    @NotNull
    private static List<String> probeForStreams(@NotNull String ffprobePath, @NotNull File file) throws IOException {
        FFMpegCommand probeCommand = new FFMpegCommand(ffprobePath);
        probeCommand.addGlobalOption("-v", "error", "-show_streams");
        probeCommand.addInputFile(file.toString());

        ProcessBuilder processBuilder = new ProcessBuilder(probeCommand.getResult());
        ProcessResult result = runProcess(processBuilder, file.getParentFile(), PROBE_SPLIT_CHARACTER, false);
        if (result.exitCode() != 0) {
            throw new IllegalArgumentException("File probe failed with code " + result.exitCode());
        }
        return StringUtil.stringBetween(result.output(), "[STREAM]", "[/STREAM]");
    }

    /**
     * Gets the duration, in seconds, of the given file
     *
     * @param ffprobePath <p>The path to the ffprobe executable</p>
     * @param file        <p>The file to get the duration of</p>
     * @return <p>The duration</p>
     * @throws IOException           <p>If unable to probe the file</p>
     * @throws NumberFormatException <p>If ffmpeg returns a non-number</p>
     */
    public static double getDuration(@NotNull String ffprobePath, @NotNull File file) throws IOException, NumberFormatException {
        FFMpegCommand probeCommand = new FFMpegCommand(ffprobePath);
        probeCommand.addGlobalOption("-v", "error", "-show_entries", "format=duration", "-of",
                "default=noprint_wrappers=1:nokey=1");
        probeCommand.addInputFile(file.toString());

        ProcessBuilder processBuilder = new ProcessBuilder(probeCommand.getResult());
        ProcessResult result = runProcess(processBuilder, file.getParentFile(), "", false);
        if (result.exitCode() != 0) {
            throw new IllegalArgumentException("File probe failed with code " + result.exitCode());
        }
        return Double.parseDouble(result.output().trim());
    }

    /**
     * Takes a list of all streams and parses each stream into one of three objects
     *
     * @param ffprobePath     <p>The path to the ffprobe executable</p>
     * @param streams         <p>A list of all streams for the current file.</p>
     * @param file            <p>The file currently being converted.</p>
     * @param subtitleFormats <p>The extensions to accept for external subtitles</p>
     * @param audioFormats    <p>The extensions to accept for external audio tracks</p>
     * @return <p>A list of StreamObjects.</p>
     */
    @NotNull
    private static StreamProbeResult parseStreams(@NotNull String ffprobePath, @NotNull List<String> streams,
                                                  @NotNull File file, @NotNull List<String> subtitleFormats,
                                                  @NotNull List<String> audioFormats) throws IOException {
        StreamProbeResult probeResult = new StreamProbeResult(new ArrayList<>(List.of(file)),
                parseStreamObjects(streams));
        getExternalStreams(probeResult, ffprobePath, file.getParentFile(), file.getName(), subtitleFormats);
        getExternalStreams(probeResult, ffprobePath, file.getParentFile(), file.getName(), audioFormats);
        return probeResult;
    }

    /**
     * Parses the stream objects found in the given streams
     *
     * @param streams <p>The stream data to parse</p>
     * @return <p>The parsed stream objects</p>
     */
    @NotNull
    private static List<StreamObject> parseStreamObjects(@NotNull List<String> streams) {
        List<StreamObject> parsedStreams = new ArrayList<>();
        int relativeAudioIndex = 0;
        int relativeVideoIndex = 0;
        int relativeSubtitleIndex = 0;

        for (String stream : streams) {
            String[] streamParts = stream.split(PROBE_SPLIT_CHARACTER);
            Map<StreamTag, String> streamInfo = getStreamInfo(streamParts);
            StreamType streamType = getStreamType(streamInfo);

            switch (streamType) {
                case VIDEO -> parsedStreams.add(new VideoStream(streamInfo, 0, relativeVideoIndex++));
                case AUDIO -> parsedStreams.add(new AudioStream(streamInfo, 0, relativeAudioIndex++));
                case SUBTITLE -> parsedStreams.add(new SubtitleStream(streamInfo, 0, relativeSubtitleIndex++));
                case OTHER -> parsedStreams.add(new OtherStream(streamInfo, 0, false));
                case COVER_IMAGE -> parsedStreams.add(new OtherStream(streamInfo, 0, true));
            }
        }
        return parsedStreams;
    }

    /**
     * Gets the type of a stream from its stream info
     *
     * @param streamInfo <p>The information describing the stream</p>
     * @return <p>The type of the stream</p>
     */
    @NotNull
    private static StreamType getStreamType(@NotNull Map<StreamTag, String> streamInfo) {
        String codecType = ValueParsingHelper.parseString(streamInfo.get(StreamTag.CODEC_TYPE), "");
        switch (codecType) {
            case "video":
                String mime = ValueParsingHelper.parseString(streamInfo.get(StreamTag.TAG_MIME_TYPE), "");
                // Some attached covers are marked as video streams
                if (ValueParsingHelper.parseInt(streamInfo.get(StreamTag.DISPOSITION_ATTACHED_PIC), 0) != 1 &&
                        !mime.startsWith("image/") && !mime.endsWith("-font")) {
                    return StreamType.VIDEO;
                } else {
                    return StreamType.COVER_IMAGE;
                }
            case "audio":
                return StreamType.AUDIO;
            case "subtitle":
                return StreamType.SUBTITLE;
            default:
                return StreamType.OTHER;
        }
    }

    /**
     * Gets stream info from the given raw stream info lines
     *
     * @param streamParts <p>The stream info lines to parse</p>
     * @return <p>The stream tag map parsed</p>
     */
    @NotNull
    private static Map<StreamTag, String> getStreamInfo(@Nullable String[] streamParts) {
        Map<StreamTag, String> streamInfo = new HashMap<>();
        for (String part : streamParts) {
            if (part == null || !part.contains("=")) {
                continue;
            }
            String[] keyValue = part.split("=");
            String value = keyValue.length > 1 ? keyValue[1] : "";

            StreamTag tag = StreamTag.getFromString(keyValue[0]);
            if (tag != null) {
                streamInfo.put(tag, value);
            }
        }
        return streamInfo;
    }

    /**
     * Tries to find any external files adjacent to the first input file, and appends it to the given probe result
     *
     * @param streamProbeResult <p>The stream probe result to append to</p>
     * @param ffprobePath       <p>The path/command to ffprobe</p>
     * @param directory         <p>The directory containing the file</p>
     * @param convertingFile    <p>The first/main file to be converted</p>
     * @param formats           <p>The extensions to accept for external tracks</p>
     */
    private static void getExternalStreams(@NotNull StreamProbeResult streamProbeResult,
                                           @NotNull String ffprobePath, @NotNull File directory,
                                           @NotNull String convertingFile,
                                           @NotNull List<String> formats) throws IOException {
        //Find all files in the same directory with external subtitle formats
        File[] files = FileUtil.listFilesRecursive(directory, formats, 1);

        //Return early if no files were found
        if (files == null) {
            return;
        }

        String fileTitle = FileUtil.stripExtension(convertingFile);
        List<File> filesList = new ArrayList<>(Arrays.asList(files));

        //Finds the files which are subtitles probably belonging to the file
        filesList = ListUtil.getMatching(filesList, (file) -> file.getName().contains(fileTitle));

        for (File file : filesList) {
            int inputIndex = streamProbeResult.parsedFiles().size();
            streamProbeResult.parsedFiles().add(file);
            //Probe the files and add them to the result list
            List<String> streams = probeForStreams(ffprobePath, file);

            int audioIndex = 0;
            int subtitleIndex = 0;
            int videoIndex = 0;
            for (String stream : streams) {
                String[] streamParts = stream.split(PROBE_SPLIT_CHARACTER);
                Map<StreamTag, String> streamInfo = getStreamInfo(streamParts);
                StreamObject streamObject = null;
                switch (getStreamType(streamInfo)) {
                    case SUBTITLE -> streamObject = new SubtitleStream(streamInfo, inputIndex, subtitleIndex++);
                    case AUDIO -> streamObject = new AudioStream(streamInfo, inputIndex, audioIndex++);
                    case VIDEO -> streamObject = new VideoStream(streamInfo, inputIndex, videoIndex++);
                }
                if (streamObject != null) {
                    streamProbeResult.parsedStreams().add(streamObject);
                }
            }
        }
    }

    /**
     * Reads from a process reader
     *
     * @param reader <p>The reader of a process.</p>
     * @return <p>The output from the readProcess.</p>
     * @throws IOException <p>On reader failure.</p>
     */
    @NotNull
    private static String readProcess(@NotNull BufferedReader reader, @NotNull String spacer) throws IOException {
        String line;
        StringBuilder text = new StringBuilder();
        while (reader.ready() && (line = reader.readLine()) != null && !line.isEmpty() && !line.equals("\n")) {
            text.append(line).append(spacer);
        }
        return text.toString().trim();
    }

    /**
     * Gets available hardware acceleration types
     *
     * @param ffmpegPath <p>The path to ffmpeg's executable</p>
     * @return <p>The available hardware acceleration methods</p>
     * @throws IOException <p>If the process fails</p>
     */
    @NotNull
    public static List<String> getHWAcceleration(@NotNull String ffmpegPath) throws IOException {
        FFMpegCommand probeCommand = new FFMpegCommand(ffmpegPath);
        probeCommand.addGlobalOption("-v", "error", "-hwaccels");

        ProcessBuilder processBuilder = new ProcessBuilder(probeCommand.getResult());
        ProcessResult result = runProcess(processBuilder, null, PROBE_SPLIT_CHARACTER, false);
        return List.of(result.output().split(PROBE_SPLIT_CHARACTER));
    }

}