Chords/src/main/java/moe/nekojimi/chords/Downloader.java

/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package moe.nekojimi.chords;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.StringReader;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import javax.json.Json;
import javax.json.JsonArray;
import javax.json.JsonObject;
import javax.json.JsonReader;
import moe.nekojimi.chords.Downloader.DownloadTask;
import net.dv8tion.jda.api.audio.AudioSendHandler;

/**
 *
 * @author jimj316
 */
public class Downloader implements Consumer<DownloadTask>
{

    private static final int DOWNLOAD_TIMEOUT = 300;
    private static final int INFO_TIMEOUT = 60;
    private static final int FORMAT_TIMEOUT = 5;

    private static final int BITRATE_TARGET = (int) AudioSendHandler.INPUT_FORMAT.getSampleRate();
    private static final Pattern FORMAT_PATTERN = Pattern.compile("^([\\w]+)\\s+([\\w]+)\\s+(\\w+ ?\\w*)\\s+(.*)$");
    public static final Pattern DESTINATION_PATTERN = Pattern.compile("Destination: (.*\\.wav)");
    public static final Pattern PROGRESS_PATTERN = Pattern.compile("\\[download\\].*?([\\d\\.]+)%");
    private static final Pattern ETA_PATTERN = Pattern.compile("\\[download\\].*?ETA\\s+(\\d{1,2}:\\d{2})");
    private static final Pattern DOWNLOAD_ITEM_PATTERN = Pattern.compile("\\[download\\] Downloading item (\\d+) of (\\d+)");

    private final List<DownloadTask> downloadQueue = new LinkedList<>();
    private final LinkedBlockingDeque<Runnable> workQueue = new LinkedBlockingDeque<>();
    private final ThreadPoolExecutor executor = new ThreadPoolExecutor(2, 4, 30, TimeUnit.SECONDS, workQueue);
    private final ScheduledExecutorService scheduler = new ScheduledThreadPoolExecutor(1);
//    private Consumer<Track> next;
    private BiConsumer<TrackRequest, Exception> messageHandler;

    private File downloadDir = null;

    private int trackNumber = 1;

    public Downloader()
    {

    }

    @Override
    public void accept(DownloadTask task)
    {
        // if all tracks of the request are already downloaded, just skip
        if (!task.request.getTracks().isEmpty() && task.request.getTracks().stream().allMatch((t) -> t.isDownloaded()))
        {
            for (Track track : task.request.getTracks())
                task.getDestination().accept(track);
            return;
        }

        downloadQueue.add(task);
        getInfo(task.request);
        // TODO: get info should also use the thread pool
        executor.submit(() ->
        {
            try
            {
//                getFormats(track);
                download(task);
            } catch (Exception ex)
            {
                ex.printStackTrace();
            }
        });
    }

    private List<Format> sortFormats(Collection<Format> input)
    {
        List<Format> formats = new ArrayList<>(input);
        formats.sort((Format a, Format b) ->
        {
            // audio only preferred to video
            int comp = 0;
            comp = Boolean.compare(a.isAudioOnly(), b.isAudioOnly());
            if (comp == 0)
            {
                // known preferred to unknown
                if (a.getSampleRate() == b.getSampleRate())
                    comp = 0;
                else if (a.getSampleRate() <= 0)
                    comp = 1;
                else if (b.getSampleRate() <= 0)
                    comp = -1;
                else // closer to the target bitrate is best
                {
                    int aDist = Math.abs(a.getSampleRate() - BITRATE_TARGET);
                    int bDist = Math.abs(b.getSampleRate() - BITRATE_TARGET);
                    comp = Integer.compare(bDist, aDist);
                }
            }
            if (comp == 0)
            {
                // known preferred to unknown
                if (a.getSize() == b.getSize())
                    comp = 0;
                else if (a.getSize() <= 0)
                    comp = 1;
                else if (b.getSize() <= 0)
                    comp = -1;
                else // smaller is better
                    comp = Long.compare(b.getSize(), a.getSize());
            }
            return -comp;
        });
        return formats;
    }

    private void getFormats(Track track)
    {
        if (!track.getFormats().isEmpty())
            return;
        try
        {
            String cmd = Chords.getSettings().getYtdlCommand() + " --skip-download -F " + track.getUrl().toString();
            Process exec = runCommand(cmd, FORMAT_TIMEOUT);
            InputStream input = exec.getInputStream();
            String output = new String(input.readAllBytes(), Charset.defaultCharset());

            List<Format> formats = new ArrayList<>();

            int codeCol = 0;
            int extCol = 1;
            int resCol = 2;
            int noteCol = 3;
            int sizeCol = -1;
            int bitrateCol = -1;

            List<String> list = output.lines().collect(Collectors.toList());
            int i = 0;
            while (!list.get(i).contains("Available formats"))
                i++;
            i++;
            if (list.get(i).contains("FILESIZE"))
            {
                String[] split = list.get(i).split("\\s+");
                for (int j = 0; j < split.length; j++)
                {
                    switch (split[j])
                    {
                        case "ID":
                            codeCol = j;
                            break;
                        case "EXT":
                            extCol = j;
                            break;
                        case "RESOLUTION":
                            resCol = j;
                            break;
                        case "MORE":
                            noteCol = j;
                            break;
                        case "FILESIZE":
                            sizeCol = j;
                            break;
                        case "TBR":
                            bitrateCol = j;
                    }
                }
                i += 2;
            }
            for (; i < list.size(); i++)
            {
                String line = list.get(i);
                String[] split = line.split("\\s+", Math.max(4, noteCol - 1));
                if (split.length < 4)
                    continue;
                final Format format = new Format(split[codeCol], split[extCol], split[resCol], split[noteCol]);
                if (sizeCol >= 0)
                    format.setSize(Util.parseSize(split[sizeCol]));
                if (bitrateCol >= 0)
                    format.setSampleRate(Util.parseSampleRate(split[bitrateCol]));
                formats.add(format);
            }
            track.setFormats(formats);

        } catch (Exception ex)
        {
            Logger.getLogger(Downloader.class.getName()).log(Level.SEVERE, null, ex);
//            return List.of();
        }
    }

    private List<Track> getInfo(TrackRequest request)
    {
        List<Track> ret = new ArrayList<>();
        try
        {
            String cmd = Chords.getSettings().getYtdlCommand() + " --skip-download --print-json " + request.getUrl().toString();
            Process exec = runCommand(cmd, INFO_TIMEOUT);
            InputStream input = exec.getInputStream();

            // read each line as JSON, turn each into a track object
            Scanner sc = new Scanner(input);
            while (sc.hasNextLine())
            {
                Track track = new Track(request.getUrl());
                track.setNumber(trackNumber);
                trackNumber++;
                request.addTrack(track);

                String line = sc.nextLine();
                JsonReader reader = Json.createReader(new StringReader(line));
                JsonObject object = reader.readObject();
                if (track.getTitle() == null)
                    track.setTitle(object.getString("title", null));
                if (track.getArtist() == null)
                    track.setArtist(object.getString("uploader", null));

                JsonArray formatsJSON = object.getJsonArray("formats");
                if (formatsJSON != null)
                {
                    List<Format> formats = new ArrayList<>();
                    for (JsonObject formatJson : formatsJSON.getValuesAs(JsonObject.class))
                    {
                        Format format = Format.fromJSON(formatJson);
                        if (format != null)
                            formats.add(format);
                    }
                    track.setFormats(formats);
                }

                ret.add(track);
            }
        } catch (Exception ex)
        {
            Logger.getLogger(Downloader.class.getName()).log(Level.SEVERE, null, ex);
        }
        return ret;
    }

    private File getDownloadDir() throws IOException
    {
        if (downloadDir == null || !downloadDir.exists() || !downloadDir.canWrite())
            downloadDir = Files.createTempDirectory("chords").toFile();
        return downloadDir;
    }

    private Track getTrackFromRequest(TrackRequest request, int idx)
    {
        // if there's less tracks in the request than expected, fill the array
        while (idx >= request.getTracks().size())
        {
            Track track = new Track(request.getUrl());
            track.setNumber(trackNumber);
            trackNumber++;
            request.addTrack(track);
        }
        return request.getTracks().get(idx);
    }

    private void download(DownloadTask task)
    {
        Set<Format> uniqueFormats = new HashSet<>();
        for (Track track : task.request.getTracks())
        {
            uniqueFormats.addAll(track.getFormats());
        }

        List<Format> sortedFormats = sortFormats(uniqueFormats);
        String formatCodes = "";
        final List<Format> formats = sortedFormats;
        for (int i = 0; i < 5 && i < sortedFormats.size(); i++)
            formatCodes += formats.get(i).getCode() + "/";

        int downloadIdx = 0;

        try
        {
            messageHandler.accept(task.request, null);
            String cmd = Chords.getSettings().getYtdlCommand()
                    + " -x"
                    + " -f " + formatCodes + "worstaudio/bestaudio/worst/best"
                    + " --audio-format=wav"
                    + " --no-playlist"
                    + " --extractor-args youtube:player_client=android"
                    + " -o " + getDownloadDir().getAbsolutePath() + "/%(title)s.%(ext)s "
                    + task.request.getUrl().toString();

            Process exec = runCommand(cmd, DOWNLOAD_TIMEOUT);
            InputStream in = exec.getInputStream();

            Scanner sc = new Scanner(in);
            while (sc.hasNextLine())
            {
                String line = sc.nextLine();
                System.out.println(line);

                Matcher itemMatcher = DOWNLOAD_ITEM_PATTERN.matcher(line);
                if (itemMatcher.find())
                {
                    int idx = Integer.parseInt(itemMatcher.group(1)) - 1;
//                    int total = Integer.parseInt(itemMatcher.group(2));

                    downloadIdx = idx;
                }

                Matcher progMatcher = PROGRESS_PATTERN.matcher(line);
                if (progMatcher.find())
                {
                    getTrackFromRequest(task.request, downloadIdx).setProgress(Double.parseDouble(progMatcher.group(1)));
                    messageHandler.accept(task.request, null);
                }

                Matcher destMatcher = DESTINATION_PATTERN.matcher(line);
                if (destMatcher.find())
                {
                    Track track = getTrackFromRequest(task.request, downloadIdx);

                    track.setLocation(new File(destMatcher.group(1)));

                    // this is currently our criteria for completion; submit the track and move on
                    if (task.getDestination() != null)
                        task.getDestination().accept(track);

                    track.setProgress(100.0);

                    messageHandler.accept(task.request, null);

                    downloadIdx++;
                }
            }

            //            String output = new String(in.readAllBytes(), Charset.defaultCharset());
            String error = new String(exec.getErrorStream().readAllBytes(), Charset.defaultCharset());
            //            System.out.println(output);

            if (exec.exitValue() != 0)
                throw new RuntimeException("youtube-dl failed with error " + exec.exitValue() + ", output:\n" + error);

//            task.request.setProgress(100);
            //            return true;

            downloadQueue.remove(task);
            messageHandler.accept(task.request, null);
        } catch (Exception ex)
        {
            Logger.getLogger(Downloader.class.getName()).log(Level.SEVERE, null, ex);
            if (messageHandler != null)
                messageHandler.accept(task.request, ex);
            downloadQueue.remove(task);
        }
    }

    private Process runCommand(String cmd, int timeoutSecs) throws RuntimeException, IOException, InterruptedException
    {
        System.out.println("Running command: " + cmd);
        //		Process exec = Runtime.getRuntime().exec().split(" "));
        Process exec = new ProcessBuilder(cmd.split(" ")).redirectOutput(ProcessBuilder.Redirect.PIPE).redirectError(ProcessBuilder.Redirect.PIPE).start();
        scheduler.schedule(() ->
        {
            if (exec.isAlive())
            {
                exec.destroyForcibly();
                System.err.println("Process " + cmd + " took too long, killing process.");
            }
        }, timeoutSecs, TimeUnit.SECONDS);
//        boolean done = exec.waitFor(timeoutSecs, TimeUnit.SECONDS);
//        if (!done)
//        {
//            exec.destroyForcibly();
//            throw new RuntimeException("Took too long, giving up.");
//        }
        return exec;
    }

    public BiConsumer<TrackRequest, Exception> getMessageHandler()
    {
        return messageHandler;
    }

    public void setMessageHandler(BiConsumer<TrackRequest, Exception> messageHandler)
    {
        this.messageHandler = messageHandler;
    }

    public List<DownloadTask> getDownloadQueue()
    {
        return downloadQueue;
    }

    public static class DownloadTask
    {

        private final TrackRequest request;
        private final Consumer<Track> destination;

        public DownloadTask(TrackRequest request, Consumer<Track> destination)
        {
            this.request = request;
            this.destination = destination;
        }

        public TrackRequest getTrack()
        {
            return request;
        }

        public Consumer<Track> getDestination()
        {
            return destination;
        }

    }

}