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.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() {

    }
    
    /*
    Developer notes:
      Setting the default track:
        -disposition:s:0 +default
        Where s,a,v is the type specifier, and the number is the relative index
        
        To unset default:
        -disposition:s:0 -default
    
      -map command for reordering:
          First number is the index of the input file, which is 0, unless more files are given as input
          Optionally, a,v,s can be used to select the type of stream to map
          Lastly, the number is either the global index of the stream, or the relative, if a type has been specified
          If only the first number is given, all streams from that file are selected
          The output file will contain all mapped streams in the order they are mapped
          
          Plan: Sort all streams by set criteria. Map all selected streams by looping using their relative index for
          selection.
       
      Streams should probably have an input index, so it's easier to treat extra subtitles more seamlessly. So, by
      including any external subtitle files as input, there would be no need to fiddle more with storing input files.
      
      Instead of storing the ffmpeg command as a list of strings, it should be stored as an object with different list
      for input arguments and output arguments. That way, it would be much easier to add input-related arguments later
      in the process.
      
      ffmpeg -codecs:
        Used for checking which codecs are available.
      
      ffmpeg -h encoder=h264_nvenc:
        Used to see available encoder presets/profiles/info
      
     */

    /**
     * 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>
     * @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) throws IOException {
        return parseStreams(ffprobePath, probeForStreams(ffprobePath, file), file, subtitleFormats);
    }

    /**
     * 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]");
    }

    /**
     * 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>
     * @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) throws IOException {
        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);

            String codecType = ValueParsingHelper.parseString(streamInfo.get(StreamTag.CODEC_TYPE), "");
            switch (codecType) {
                case "video":
                    // Some attached covers are marked as video streams
                    if (ValueParsingHelper.parseInt(streamInfo.get(StreamTag.DISPOSITION_ATTACHED_PIC), 0) != 1) {
                        parsedStreams.add(new VideoStream(streamInfo, 0, relativeVideoIndex++));
                    } else {
                        parsedStreams.add(new OtherStream(streamInfo, 0));
                    }
                    break;
                case "audio":
                    parsedStreams.add(new AudioStream(streamInfo, 0, relativeAudioIndex++));
                    break;
                case "subtitle":
                    parsedStreams.add(new SubtitleStream(streamInfo, 0, relativeSubtitleIndex++));
                    break;
                default:
                    parsedStreams.add(new OtherStream(streamInfo, 0));
            }
        }
        StreamProbeResult probeResult = new StreamProbeResult(List.of(file), parsedStreams);
        getExternalSubtitles(probeResult, ffprobePath, file.getParentFile(), file.getName(), subtitleFormats);
        return probeResult;
    }

    /**
     * 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("=");
            StreamTag tag = StreamTag.getFromString(keyValue[0]);
            if (tag != null) {
                streamInfo.put(tag, keyValue[1]);
            }
        }
        return streamInfo;
    }

    /**
     * Tries to find any external subtitles 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 subtitleFormats   <p>The extensions to accept for external subtitles</p>
     */
    private static void getExternalSubtitles(@NotNull StreamProbeResult streamProbeResult,
                                             @NotNull String ffprobePath, @NotNull File directory,
                                             @NotNull String convertingFile, @NotNull List<String> subtitleFormats) throws IOException {
        //Find all files in the same directory with external subtitle formats
        File[] subtitleFiles = FileUtil.listFilesRecursive(directory, subtitleFormats, 1);
        // TODO: Generalize this for external audio tracks

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

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

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

        for (File subtitleFile : subtitleFilesList) {
            int inputIndex = streamProbeResult.parsedFiles().size();
            streamProbeResult.parsedFiles().add(subtitleFile);
            //Probe the files and add them to the result list
            List<String> streams = probeForStreams(ffprobePath, subtitleFile);
            int relativeIndex = 0;
            for (String stream : streams) {
                String[] streamParts = stream.split(PROBE_SPLIT_CHARACTER);
                Map<StreamTag, String> streamInfo = getStreamInfo(streamParts);
                streamProbeResult.parsedStreams().add(new SubtitleStream(streamInfo, inputIndex, relativeIndex++));
            }
        }
    }

    /**
     * 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));
    }

}