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

package net.knarcraft.ffmpegconverter.converter;

import net.knarcraft.ffmpegconverter.FFMpegConvert;
import net.knarcraft.ffmpegconverter.config.Configuration;
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.MapAllModule;
import net.knarcraft.ffmpegconverter.converter.module.mapping.SetDefaultStreamModule;
import net.knarcraft.ffmpegconverter.converter.module.output.CopyAudioModule;
import net.knarcraft.ffmpegconverter.converter.module.output.CopySubtitlesModule;
import net.knarcraft.ffmpegconverter.converter.module.output.CopyVideoModule;
import net.knarcraft.ffmpegconverter.converter.module.output.FastStartModule;
import net.knarcraft.ffmpegconverter.converter.module.output.SetOutputFileModule;
import net.knarcraft.ffmpegconverter.converter.module.output.SetStreamLanguageModule;
import net.knarcraft.ffmpegconverter.converter.sorter.AudioLanguageSorter;
import net.knarcraft.ffmpegconverter.converter.sorter.MinimalSubtitleSorter;
import net.knarcraft.ffmpegconverter.converter.sorter.SpecialAudioSorter;
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 net.knarcraft.ffmpegconverter.utility.OutputUtil;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

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

/**
 * A converter for converting anime, keeping all streams
 */
public class StreamOrderConverter extends AbstractConverter {

    private final List<String> audioLanguages;
    private final List<String> subtitleLanguages;
    private final MinimalSubtitlePreference subtitlePreference;
    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>The language priority for languages</p>
     * @param subtitleLanguages <p>The language priority for subtitles</p>
     */
    public StreamOrderConverter(@NotNull String ffprobePath, @NotNull String ffmpegPath,
                                @NotNull String audioLanguages, @NotNull String subtitleLanguages,
                                @NotNull String subtitleNameFilter) {
        super("mkv");
        Configuration configuration = FFMpegConvert.getConfiguration();
        this.ffprobePath = ffprobePath;
        this.ffmpegPath = ffmpegPath;
        this.audioLanguages = List.of(audioLanguages.split(","));
        this.subtitleLanguages = List.of(subtitleLanguages.split(","));
        this.subtitlePreference = configuration.getMinimalSubtitlePreference();
        this.subtitleNameFilter = subtitleNameFilter;
    }

    @Override
    @Nullable
    public FFMpegCommand generateConversionCommand(@NotNull String executable,
                                                   @NotNull StreamProbeResult probeResult,
                                                   @NotNull String outFile) {
        FFMpegCommand command = FFMpegHelper.getFFMpegGeneralFileCommand(executable, probeResult.parsedFiles());
        List<ConverterModule> modules = new ArrayList<>();
        if (this.debug) {
            modules.add(new DebugModule(90, 120));
        }
        modules.add(new FastStartModule());

        //Map all video streams
        List<VideoStream> videoStreams = probeResult.getVideoStreams();
        modules.add(new MapAllModule<>(videoStreams));
        setOutputIndexes(videoStreams);

        //Get the first audio stream in accordance with chosen languages
        StreamSorter<AudioStream> audioSorter = new AudioLanguageSorter(this.audioLanguages)
                .append(new SpecialAudioSorter(MinimalSubtitlePreference.REJECT));
        List<AudioStream> sortedAudio = audioSorter.chainSort(probeResult.getAudioStreams());
        setOutputIndexes(sortedAudio);
        modules.add(new MapAllModule<>(sortedAudio));
        modules.add(new SetDefaultStreamModule<>(sortedAudio, 0));
        modules.add(new SetStreamLanguageModule<>(sortedAudio));

        //Get the first subtitle stream in accordance with chosen languages and signs and songs prevention
        StreamSorter<SubtitleStream> subtitleSorter = new SubtitleTitleSorter(
                List.of(this.subtitleNameFilter.split(",")))
                .append(new SubtitleLanguageSorter(this.subtitleLanguages))
                .append(new MinimalSubtitleSorter(this.subtitlePreference));
        List<SubtitleStream> sortedSubtitles = subtitleSorter.chainSort(probeResult.getSubtitleStreams());
        setOutputIndexes(sortedSubtitles);
        modules.add(new MapAllModule<>(sortedSubtitles));
        modules.add(new SetDefaultStreamModule<>(sortedSubtitles, 0));
        modules.add(new SetStreamLanguageModule<>(sortedSubtitles));
        OutputUtil.printDebug("Subtitle preference: " + this.subtitleLanguages + ", Subtitle name filter: " +
                this.subtitleNameFilter + ", Minimal Subtitle Preference: " + this.subtitlePreference.name() +
                ", Sorted order: " + Arrays.toString(sortedSubtitles.toArray()));

        for (AudioStream audioStream : sortedAudio) {
            modules.add(new CopyAudioModule(audioStream));
        }
        modules.add(new CopySubtitlesModule());
        modules.add(new CopyVideoModule());

        // Map all attached streams, such as fonts and covers
        modules.add(new MapAllModule<>(probeResult.getOtherStreams()));

        modules.add(new SetOutputFileModule(outFile));

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

    @Override
    @NotNull
    public List<String> getValidFormats() {
        return this.videoFormats;
    }

}