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

package net.knarcraft.ffmpegconverter.utility;

import net.knarcraft.ffmpegconverter.streams.AudioStream;
import net.knarcraft.ffmpegconverter.streams.StreamObject;
import net.knarcraft.ffmpegconverter.streams.SubtitleStream;
import net.knarcraft.ffmpegconverter.streams.VideoStream;

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.List;

/**
 * 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>
     * @return <p>A list of StreamObjects.</p>
     * @throws IOException <p>If the process can't be readProcess.</p>
     */
    public static List<StreamObject> probeFile(String ffprobePath, 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 fileName   <p>The name of the file to execute on.</p>
     * @return <p>A base list of ffmpeg commands for converting a video for web</p>
     */
    public static List<String> getFFMpegWebVideoCommand(String executable, String fileName) {
        List<String> command = getFFMpegGeneralFileCommand(executable, fileName);
        command.add("-vcodec");
        command.add("h264");
        command.add("-pix_fmt");
        command.add("yuv420p");
        command.add("-ar");
        command.add("48000");
        command.add("-movflags");
        command.add("+faststart");
        return command;
    }

    /**
     * Creates a list containing command line arguments for a general file
     *
     * @param executable <p>The executable to use (ffmpeg/ffprobe).</p>
     * @param fileName   <p>The name of the file to execute on.</p>
     * @return <p>A base list of ffmpeg commands for converting a file.</p>
     */
    public static List<String> getFFMpegGeneralFileCommand(String executable, String fileName) {
        List<String> command = new ArrayList<>();
        command.add(executable);
        command.add("-nostdin");
        command.add("-i");
        command.add(fileName);
        return command;
    }

    /**
     * Adds debugging parameters for only converting parts of a file
     *
     * @param command <p>The list containing the command to run.</p>
     * @param start   <p>The offset before converting.</p>
     * @param length  <p>The offset for stopping the conversion.</p>
     */
    public static void addDebugArguments(List<String> command, int start, int length) {
        command.add("-ss");
        command.add("" + start);
        command.add("-t");
        command.add("" + length);
    }

    /**
     * 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>
     */
    public static String runProcess(ProcessBuilder processBuilder, File folder, 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.equals("")) {
                if (write) {
                    OutputUtil.println(read);
                } else {
                    OutputUtil.printDebug(read);
                    output.append(read);
                }
            }
        }
        OutputUtil.println("Process finished.");
        return output.toString();
    }

    /**
     * Adds audio to a command
     *
     * @param command     <p>The command to add audio to.</p>
     * @param audioStream <p>The audio stream to be added.</p>
     * @param toStereo    <p>Whether to convert the audio stream to stereo.</p>
     */
    public static void addAudioStreams(List<String> command, AudioStream audioStream, boolean toStereo) {
        if (audioStream != null) {
            command.add("-map");
            command.add("0:" + audioStream.getAbsoluteIndex());
            if (toStereo && audioStream.getChannels() > 2) {
                command.add("-af");
                command.add("pan=stereo|FL=FC+0.30*FL+0.30*BL|FR=FC+0.30*FR+0.30*BR");
                //command.add("pan=stereo|FL < 1.0*FL + 0.707*FC + 0.707*BL|FR < 1.0*FR + 0.707*FC + 0.707*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>
     * @param file           <p>The file to convert.</p>
     */
    public static void addSubtitlesAndVideo(List<String> command, SubtitleStream subtitleStream,
                                            VideoStream videoStream, File file) {
        //No appropriate subtitle was found. Just add the video stream.
        if (subtitleStream == null) {
            command.add("-map");
            command.add(String.format("0:%d", videoStream.getAbsoluteIndex()));
            return;
        }

        //Add the correct command arguments depending on the subtitle type
        if (!subtitleStream.getIsImageSubtitle()) {
            addSubtitle(command, subtitleStream, videoStream);
        } else if (file.getName().equals(subtitleStream.getFile())) {
            addInternalImageSubtitle(command, subtitleStream, videoStream);
        } else {
            addExternalImageSubtitle(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 addSubtitle(List<String> command, SubtitleStream subtitleStream, VideoStream videoStream) {
        command.add("-map");
        command.add(String.format("0:%d", videoStream.getAbsoluteIndex()));
        command.add("-vf");
        String safeFileName = escapeSpecialCharactersInFileName(subtitleStream.getFile());
        String subtitleCommand = String.format("subtitles='%s':si=%d", safeFileName,
                subtitleStream.getRelativeIndex());
        command.add(subtitleCommand);
    }

    /**
     * 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 image 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 addInternalImageSubtitle(List<String> command, SubtitleStream subtitleStream,
                                                 VideoStream videoStream) {
        command.add("-filter_complex");
        String filter = String.format("[0:v:%d][0:%d]overlay", videoStream.getAbsoluteIndex(),
                subtitleStream.getAbsoluteIndex());
        command.add(filter);
    }

    /**
     * Adds external image subtitle commands to a command list
     *
     * @param command               <p>The list containing the FFmpeg commands.</p>
     * @param externalImageSubtitle <p>The external image subtitle stream to add.</p>
     * @param videoStream           <p>The video stream to burn the subtitle into.</p>
     */
    private static void addExternalImageSubtitle(List<String> command, SubtitleStream externalImageSubtitle,
                                                 VideoStream videoStream) {
        command.add("-i");
        command.add(externalImageSubtitle.getFile());
        command.add("-filter_complex");
        command.add(String.format("[1:s]scale=width=%d:height=%d,crop=w=%d:h=%d:x=0:y=out_h[sub];[%d:v]" +
                        "[sub]overlay", videoStream.getWidth(), videoStream.getHeight(), videoStream.getWidth(),
                videoStream.getHeight(), videoStream.getAbsoluteIndex()));
        command.add("-profile:v");
        command.add("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 {
        ProcessBuilder processBuilder = new ProcessBuilder(
                ffprobePath,
                "-v",
                "error",
                "-show_entries",
                "stream_tags=language,title:stream=index,codec_name,codec_type,channels,codec_type,width,height",
                file.toString()
        );
        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>
     */
    private static List<StreamObject> parseStreams(String ffprobePath, String[] streams, 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);
            if (stream.contains("codec_type=video")) {
                parsedStreams.add(parseVideoStream(streamParts, relativeVideoIndex++));
            } else if (stream.contains("codec_type=audio")) {
                parsedStreams.add(parseAudioStream(streamParts, relativeAudioIndex++));
            } else if (stream.contains("codec_type=subtitle")) {
                parsedStreams.add(parseSubtitleStream(streamParts, relativeSubtitleIndex++, file.getName()));
            }
        }
        List<StreamObject> externalSubtitles = getExternalSubtitles(ffprobePath, file.getParentFile(), file.getName());
        parsedStreams.addAll(externalSubtitles);
        return parsedStreams;
    }

    /**
     * Checks whether there exists an external image subtitle with the same filename as the file
     *
     * @param ffprobePath    <p>The path/command to ffprobe.</p>
     * @param directory      <p>The directory containing the file.</p>
     * @param convertingFile <p>The file to be converted.</p>
     * @return <p>The extension of the subtitle or empty if no subtitle was found.</p>
     */
    private static List<StreamObject> getExternalSubtitles(String ffprobePath, File directory, String convertingFile)
            throws IOException {
        List<StreamObject> parsedStreams = new ArrayList<>();
        //Find all files in the same directory with external subtitle formats
        String[] subtitleFormats = FileUtil.readFileLines("subtitle_formats.txt");
        File[] subtitleFiles = FileUtil.listFilesRecursive(directory, subtitleFormats, 1);

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

        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) {
            //Probe the files and add them to the result list
            String[] streams = probeForStreams(ffprobePath, subtitleFile);
            for (String stream : streams) {
                String[] streamParts = stream.split(PROBE_SPLIT_CHARACTER);
                parsedStreams.add(parseSubtitleStream(streamParts, 0, subtitleFile.getName()));
            }
        }
        return parsedStreams;
    }

    /**
     * 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.equals("") && !line.equals("\n")) {
            text.append(line).append(spacer);
        }
        return text.toString().trim();
    }

    /**
     * Parses a list of video stream parameters to a video stream object
     *
     * @param streamParts   <p>A list of parameters belonging to an video stream.</p>
     * @param relativeIndex <p>The relative index of the video stream.</p>
     * @return <p>A SubtitleStream object.</p>
     * @throws NumberFormatException <p>If codec index contains a non-numeric value.</p>
     */
    private static VideoStream parseVideoStream(String[] streamParts, int relativeIndex) throws NumberFormatException {
        String codec = null;
        int absoluteIndex = -1;
        int width = -1;
        int height = -1;
        for (String streamPart : streamParts) {
            if (streamPart.startsWith("codec_name=")) {
                codec = streamPart.replace("codec_name=", "");
            } else if (streamPart.startsWith("index=")) {
                absoluteIndex = Integer.parseInt(streamPart.replace("index=", ""));
            } else if (streamPart.startsWith("width=")) {
                width = Integer.parseInt(streamPart.replace("width=", ""));
            } else if (streamPart.startsWith("height=")) {
                height = Integer.parseInt(streamPart.replace("height=", ""));
            }
        }
        return new VideoStream(codec, absoluteIndex, relativeIndex, width, height);
    }

    /**
     * Parses a list of audio stream parameters to an audio stream object
     *
     * @param streamParts   <p>A list of parameters belonging to an audio stream.</p>
     * @param relativeIndex <p>The relative index of the audio stream.</p>
     * @return <p>A SubtitleStream object.</p>
     * @throws NumberFormatException <p>If codec index contains a non-numeric value.</p>
     */
    private static AudioStream parseAudioStream(String[] streamParts, int relativeIndex) throws NumberFormatException {
        String codec = null;
        int absoluteIndex = -1;
        String language = null;
        int channels = 0;
        String title = "";
        for (String streamPart : streamParts) {
            if (streamPart.startsWith("codec_name=")) {
                codec = streamPart.replace("codec_name=", "");
            } else if (streamPart.startsWith("index=")) {
                absoluteIndex = Integer.parseInt(streamPart.replace("index=", ""));
            } else if (streamPart.startsWith("TAG:language=")) {
                language = streamPart.replace("TAG:language=", "");
            } else if (streamPart.startsWith("channels=")) {
                channels = Integer.parseInt(streamPart.replace("channels=", ""));
            } else if (streamPart.startsWith("TAG:title=")) {
                title = streamPart.replace("TAG:title=", "");
            }
        }
        return new AudioStream(codec, absoluteIndex, relativeIndex, language, title, channels);
    }

    /**
     * Parses a list of subtitle stream parameters to a subtitle stream object
     *
     * @param streamParts   <p>A list of parameters belonging to a subtitle stream.</p>
     * @param relativeIndex <p>The relative index of the subtitle.</p>
     * @param file          <p>The file currently being converted.</p>
     * @return <p>A SubtitleStream object.</p>
     * @throws NumberFormatException <p>If codec index contains a non-numeric value.</p>
     */
    private static SubtitleStream parseSubtitleStream(String[] streamParts, int relativeIndex, String file)
            throws NumberFormatException {
        String codecName = null;
        int absoluteIndex = -1;
        String language = null;
        String title = "";
        for (String streamPart : streamParts) {
            if (streamPart.startsWith("codec_name=")) {
                codecName = streamPart.replace("codec_name=", "");
            } else if (streamPart.startsWith("index=")) {
                absoluteIndex = Integer.parseInt(streamPart.replace("index=", ""));
            } else if (streamPart.startsWith("TAG:language=")) {
                language = streamPart.replace("TAG:language=", "");
            } else if (streamPart.startsWith("TAG:title=")) {
                title = streamPart.replace("TAG:title=", "");
            }
        }
        return new SubtitleStream(codecName, absoluteIndex, relativeIndex, language, title, file);
    }
}