FFmpegConvert/src/main/java/net/knarcraft/ffmpegconverter/converter/WebAnimeConverter.java

package net.knarcraft.ffmpegconverter.converter;

import net.knarcraft.ffmpegconverter.container.FFMpegCommand;
import net.knarcraft.ffmpegconverter.container.StreamProbeResult;
import net.knarcraft.ffmpegconverter.converter.module.ConverterModule;
import net.knarcraft.ffmpegconverter.converter.module.DebugModule;
import net.knarcraft.ffmpegconverter.converter.module.ModuleExecutor;
import net.knarcraft.ffmpegconverter.converter.module.mapping.AddStereoAudioStreamModule;
import net.knarcraft.ffmpegconverter.converter.module.mapping.BurnSubtitleModule;
import net.knarcraft.ffmpegconverter.converter.module.mapping.SelectSingleStreamModule;
import net.knarcraft.ffmpegconverter.converter.sorter.AudioLanguageSorter;
import net.knarcraft.ffmpegconverter.converter.sorter.MinimalSubtitleSorter;
import net.knarcraft.ffmpegconverter.converter.sorter.StreamSorter;
import net.knarcraft.ffmpegconverter.converter.sorter.SubtitleLanguageSorter;
import net.knarcraft.ffmpegconverter.converter.sorter.SubtitleTitleSorter;
import net.knarcraft.ffmpegconverter.property.MinimalSubtitlePreference;
import net.knarcraft.ffmpegconverter.streams.AudioStream;
import net.knarcraft.ffmpegconverter.streams.SubtitleStream;
import net.knarcraft.ffmpegconverter.streams.VideoStream;
import net.knarcraft.ffmpegconverter.utility.FFMpegHelper;
import org.jetbrains.annotations.NotNull;

import java.util.ArrayList;
import java.util.List;

import static net.knarcraft.ffmpegconverter.utility.FFMpegHelper.getNthSteam;

/**
 * A converter mainly designed for converting anime to web-playable mp4
 */
public class WebAnimeConverter extends AbstractConverter {

    private final String[] audioLanguages;
    private final String[] subtitleLanguages;
    private final boolean toStereo;
    private final MinimalSubtitlePreference subtitlePreference;
    private final int forcedAudioIndex;
    private final int forcedSubtitleIndex;
    private final String subtitleNameFilter;

    /**
     * Instantiates a new anime converter
     *
     * @param ffprobePath         <p>Path/command to ffprobe.</p>
     * @param ffmpegPath          <p>Path/command to ffmpeg.</p>
     * @param audioLanguages      <p>List of wanted audio languages in descending order.</p>
     * @param subtitleLanguages   <p>List of wanted subtitle languages in descending order.</p>
     * @param toStereo            <p>Convert video with several audio channels to stereo.</p>
     * @param subtitlePreference  <p>How minimal subtitles should be prioritized</p>
     * @param forcedAudioIndex    <p>A specific audio stream to force. 0-indexed from the first audio stream found</p>
     * @param forcedSubtitleIndex <p>A specific subtitle stream to force. 0-indexed for the first subtitle stream found</p>
     */
    public WebAnimeConverter(@NotNull String ffprobePath, @NotNull String ffmpegPath, @NotNull String[] audioLanguages,
                             @NotNull String[] subtitleLanguages, boolean toStereo,
                             @NotNull MinimalSubtitlePreference subtitlePreference, int forcedAudioIndex,
                             int forcedSubtitleIndex, @NotNull String subtitleNameFilter) {
        super("mp4");
        this.ffprobePath = ffprobePath;
        this.ffmpegPath = ffmpegPath;
        this.audioLanguages = audioLanguages;
        this.subtitleLanguages = subtitleLanguages;
        this.toStereo = toStereo;
        this.subtitlePreference = subtitlePreference;
        this.forcedAudioIndex = forcedAudioIndex;
        this.forcedSubtitleIndex = forcedSubtitleIndex;
        this.subtitleNameFilter = subtitleNameFilter;
    }

    @Override
    public String[] generateConversionCommand(@NotNull String executable, @NotNull StreamProbeResult probeResult,
                                              @NotNull String outFile) {
        FFMpegCommand command = FFMpegHelper.getFFMpegWebVideoCommand(executable, probeResult.parsedFiles());
        List<ConverterModule> modules = new ArrayList<>();
        if (this.debug) {
            modules.add(new DebugModule(0, 0));
            FFMpegHelper.addDebugArguments(command, 50, 120);
        }

        //Get the first audio stream in accordance with chosen languages
        AudioStream audioStream = getNthSteam(new AudioLanguageSorter(this.audioLanguages).sort(
                probeResult.getAudioStreams()), this.forcedAudioIndex);
        if (audioStream == null) {
            throw new IllegalArgumentException("The given input resulted in no audio stream being selected");
        }
        if (this.toStereo) {
            modules.add(new AddStereoAudioStreamModule(audioStream, true));
        } else {
            modules.add(new SelectSingleStreamModule(audioStream));
        }

        //Get the first video stream
        VideoStream videoStream = getNthSteam(probeResult.getVideoStreams(), 0);

        //Get the first subtitle stream in accordance with chosen languages and signs and songs prevention
        StreamSorter<SubtitleStream> subtitleSorter = new SubtitleTitleSorter(this.subtitleNameFilter)
                .append(new MinimalSubtitleSorter(this.subtitlePreference))
                .append(new SubtitleLanguageSorter(this.subtitleLanguages));
        SubtitleStream subtitleStream = getNthSteam(subtitleSorter.chainSort(probeResult.getSubtitleStreams()),
                this.forcedSubtitleIndex);

        if (subtitleStream != null && videoStream != null) {
            modules.add(new BurnSubtitleModule(subtitleStream, videoStream, true));
        } else if (videoStream != null) {
            modules.add(new SelectSingleStreamModule(videoStream));
        } else {
            throw new IllegalArgumentException("The selected video stream does not exist!");
        }

        new ModuleExecutor(command, modules).execute();
        command.setOutputFile(outFile);
        return command.getResult();
    }

    @Override
    public String[] getValidFormats() {
        return this.videoFormats;
    }

}