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

package net.knarcraft.ffmpegconverter.converter;

import net.knarcraft.ffmpegconverter.FFMpegConvert;
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.hardwarecoding.H265HardwareEncodingModule;
import net.knarcraft.ffmpegconverter.converter.module.hardwarecoding.HardwareDecodeModule;
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.SetQualityModule;
import net.knarcraft.ffmpegconverter.converter.module.output.SetVideoCodecModule;
import net.knarcraft.ffmpegconverter.converter.sorter.AudioLanguageSorter;
import net.knarcraft.ffmpegconverter.converter.sorter.ForcedFirstSorter;
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 org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

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

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

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

    /**
     * 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 subtitlePreference  <p>How minimal subtitles should be prioritized</p>
     * @param forcedAudioIndex    <p>A specific audio stream to force as default. 0-indexed from the first audio stream found</p>
     * @param forcedSubtitleIndex <p>A specific subtitle stream to force as default. 0-indexed for the first subtitle stream found</p>
     * @param forceVideoEncoding  <p>Whether to enforce encoding on the video, even if already hevc</p>
     * @param forceAudioEncoding  <p>Whether to convert audio to web-playable, even though there should be no need</p>
     */
    public AnimeConverter(@NotNull String ffprobePath, @NotNull String ffmpegPath, @NotNull String[] audioLanguages,
                          @NotNull String[] subtitleLanguages,
                          @NotNull MinimalSubtitlePreference subtitlePreference, int forcedAudioIndex,
                          int forcedSubtitleIndex, @NotNull String subtitleNameFilter, boolean forceVideoEncoding,
                          boolean forceAudioEncoding) {
        super("mkv");
        this.ffprobePath = ffprobePath;
        this.ffmpegPath = ffmpegPath;
        this.audioLanguages = audioLanguages;
        this.subtitleLanguages = subtitleLanguages;
        this.subtitlePreference = subtitlePreference;
        this.forcedAudioIndex = forcedAudioIndex;
        this.forcedSubtitleIndex = forcedSubtitleIndex;
        this.subtitleNameFilter = subtitleNameFilter;
        this.forceVideoEncoding = forceVideoEncoding;
        this.forceAudioEncoding = forceAudioEncoding;
    }

    @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());

        //Get the first video stream
        modules.add(new MapAllModule<>(probeResult.getVideoStreams()));

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

        //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))
                .append(new ForcedFirstSorter<>(this.forcedSubtitleIndex));
        List<SubtitleStream> sorted = subtitleSorter.chainSort(probeResult.getSubtitleStreams());
        modules.add(new MapAllModule<>(sorted));
        modules.add(new SetDefaultStreamModule<>(sorted, 0));

        modules.add(new HardwareDecodeModule());
        if (!this.forceAudioEncoding) {
            modules.add(new CopyAudioModule());
        }
        modules.add(new CopySubtitlesModule());

        boolean encodingNecessary = false;
        for (VideoStream videoStream : probeResult.getVideoStreams()) {
            if (!videoStream.getCodecName().trim().equalsIgnoreCase("hevc")) {
                encodingNecessary = true;
                break;
            }
        }

        if (encodingNecessary || this.forceVideoEncoding) {
            List<String> availableHWAcceleration = getAvailableHardwareEncodingMethods();
            if (FFMpegConvert.useHardwareAcceleration() && availableHWAcceleration.contains("qsv")) {
                modules.add(new SetVideoCodecModule("hevc_qsv"));
                modules.add(new SetQualityModule(17, "veryslow"));
            } else if (FFMpegConvert.useHardwareAcceleration() && availableHWAcceleration.contains("cuda")) {
                modules.add(new H265HardwareEncodingModule(20));
            } else {
                modules.add(new SetVideoCodecModule("hevc"));
                modules.add(new SetQualityModule(19, "medium"));
            }
        } else {
            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;
    }

}