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.StreamProbeResult;
import net.knarcraft.ffmpegconverter.streams.AudioStream;
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 static String[] subtitleFormats = null;

    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.
     */

    /**
     * Gets streams from a file
     *
     * @param ffprobePath <p>The path/command to ffprobe.</p>
     * @param file        <p>The file to probe.</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) throws IOException {
        return parseStreams(ffprobePath, probeForStreams(ffprobePath, file), file);
    }

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

    /**
     * Adds debugging parameters for only converting parts of a file
     *
     * @param command  <p>The command to add to</p>
     * @param start    <p>The time to start at</p>
     * @param duration <p>The duration of video to output</p>
     */
    public static void addDebugArguments(@NotNull FFMpegCommand command, int start, int duration) {
        command.addInputFileOption("-ss", String.valueOf(start));
        command.addOutputFileOption("-t", String.valueOf(duration));
    }

    /**
     * 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>
     * @throws IOException <p>If the process can't be readProcess</p>
     */
    @NotNull
    public static String runProcess(@NotNull ProcessBuilder processBuilder, @NotNull 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
        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);
            }
        }
        OutputUtil.println("Process finished.");
        return output.toString();
    }

    /**
     * Maps an audio track to a ffmpeg command's output
     *
     * @param command     <p>The command to add the audio track to</p>
     * @param audioStream <p>The audio stream to be mapped</p>
     * @param toStereo    <p>Whether to convert the audio stream to stereo</p>
     */
    public static void addAudioStream(@NotNull FFMpegCommand command, @NotNull AudioStream audioStream, boolean toStereo) {
        mapStream(command, audioStream);
        if (toStereo && audioStream.getChannels() > 2) {
            command.addOutputFileOption("-af", "pan=stereo|FL=FC+0.30*FL+0.30*BL|FR=FC+0.30*FR+0.30*BR");
        }
    }

    /**
     * Adds subtitles and video mapping to a command
     *
     * @param command        <p>The list containing the rest of the command.</p>
     * @param subtitleStream <p>The subtitle stream to be used.</p>
     * @param videoStream    <p>The video stream to be used.</p>
     */
    public static void addSubtitleAndVideoStream(@NotNull FFMpegCommand command, @Nullable SubtitleStream subtitleStream,
                                                 @NotNull VideoStream videoStream) {
        //No appropriate subtitle was found. Just add the video stream.
        if (subtitleStream == null) {
            mapStream(command, videoStream);
            return;
        }

        //Add the correct command arguments depending on the subtitle type
        if (!subtitleStream.getIsImageSubtitle()) {
            addBurnedInSubtitle(command, subtitleStream, videoStream);
        } else {
            addBurnedInImageSubtitle(command, subtitleStream, videoStream);
        }
    }

    /**
     * Adds subtitle commands to a command list
     *
     * @param command        <p>The list containing the FFmpeg commands.</p>
     * @param subtitleStream <p>The subtitle stream to add.</p>
     * @param videoStream    <p>The video stream to burn the subtitle into.</p>
     */
    private static void addBurnedInSubtitle(@NotNull FFMpegCommand command, @NotNull SubtitleStream subtitleStream,
                                            @NotNull VideoStream videoStream) {
        mapStream(command, videoStream);

        String safeFileName = escapeSpecialCharactersInFileName(
                command.getInputFiles().get(subtitleStream.getInputIndex()));
        String subtitleCommand = String.format("subtitles='%s':si=%d", safeFileName,
                subtitleStream.getRelativeIndex());
        command.addOutputFileOption("-vf", subtitleCommand);
    }

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

    /**
     * Adds arguments for using hardware acceleration during h264 decoding
     *
     * @param command <p>The command to add the arguments to</p>
     */
    public static void addH26xHardwareDecoding(@NotNull FFMpegCommand command) {
        command.addInputFileOption("-hwaccel", "cuda");
        command.addInputFileOption("-hwaccel_output_format", "cuda");
    }

    /**
     * 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>
     */
    public static void mapAllStreams(@NotNull FFMpegCommand command, @NotNull List<StreamObject> 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>
     */
    private static String escapeSpecialCharactersInFileName(String fileName) {
        return fileName.replaceAll("\\\\", "\\\\\\\\\\\\\\\\")
                .replaceAll("'", "'\\\\\\\\\\\\\''")
                .replaceAll("%", "\\\\\\\\\\\\%")
                .replaceAll(":", "\\\\\\\\\\\\:")
                .replace("]", "\\]")
                .replace("[", "\\[");
    }

    /**
     * Adds external image subtitle commands to a command list
     *
     * @param command        <p>The FFMPEG command to modify</p>
     * @param subtitleStream <p>The external image subtitle stream to burn in</p>
     * @param videoStream    <p>The video stream to burn the subtitle into</p>
     */
    private static void addBurnedInImageSubtitle(@NotNull FFMpegCommand command,
                                                 @NotNull SubtitleStream subtitleStream,
                                                 @NotNull VideoStream videoStream) {
        command.addOutputFileOption("-filter_complex",
                String.format("[%d:%d]scale=width=%d:height=%d,crop=w=%d:h=%d:x=0:y=out_h[sub];[%d:%d][sub]overlay",
                        subtitleStream.getInputIndex(), subtitleStream.getAbsoluteIndex(), videoStream.getWidth(),
                        videoStream.getHeight(), videoStream.getWidth(), videoStream.getHeight(),
                        videoStream.getInputIndex(), videoStream.getAbsoluteIndex()));
        command.addOutputFileOption("-profile:v", "baseline");
    }

    /**
     * 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>
     */
    private static String[] probeForStreams(String ffprobePath, 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());
        String result = runProcess(processBuilder, file.getParentFile(), PROBE_SPLIT_CHARACTER, false);
        return StringUtil.stringBetween(result, "[STREAM]", "[/STREAM]");
    }

    /**
     * Takes a list of all streams and parses each stream into one of three objects
     *
     * @param streams <p>A list of all streams for the current file.</p>
     * @param file    <p>The file currently being converted.</p>
     * @return <p>A list of StreamObjects.</p>
     */
    @NotNull
    private static StreamProbeResult parseStreams(@NotNull String ffprobePath, @NotNull String[] streams,
                                                  @NotNull File file) 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":
                    parsedStreams.add(new VideoStream(streamInfo, 0, relativeVideoIndex++));
                    break;
                case "audio":
                    parsedStreams.add(new AudioStream(streamInfo, 0, relativeAudioIndex++));
                    break;
                case "subtitle":
                    parsedStreams.add(new SubtitleStream(streamInfo, 0, relativeSubtitleIndex++));
                    break;
            }
        }
        StreamProbeResult probeResult = new StreamProbeResult(List.of(file), parsedStreams);
        getExternalSubtitles(probeResult, ffprobePath, file.getParentFile(), file.getName());
        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 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>
     */
    private static void getExternalSubtitles(@NotNull StreamProbeResult streamProbeResult,
                                             @NotNull String ffprobePath, @NotNull File directory,
                                             @NotNull String convertingFile) throws IOException {
        //Find all files in the same directory with external subtitle formats
        if (subtitleFormats == null) {
            subtitleFormats = FileUtil.readFileLines("subtitle_formats.txt");
        }
        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
            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>
     */
    private static String readProcess(BufferedReader reader, 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();
    }

}