Discord bot that plays music from every website ever via youtube-dl
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Chords/src/main/java/moe/nekojimi/chords/Downloader.java

302 lines
10 KiB

/*
* 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.nio.charset.Charset;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
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 = 10;
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)");
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 Consumer<Song> next;
private BiConsumer<SongRequest, Exception> messageHandler;
private File downloadDir = null;
public Downloader()
{
}
@Override
public void accept(DownloadTask task)
{
// if already downloaded, just skip
Song song = task.request.getSong();
if (song.isDownloaded())
{
task.getDestination().accept(song);
return;
}
downloadQueue.add(task);
getInfo(song);
executor.submit(() ->
{
try
{
getFormats(song);
download(task);
} catch (Exception ex)
{
ex.printStackTrace();
}
});
}
private void chooseFormats(Song song)
{
List<Format> formats = song.getFormats();
if (formats.isEmpty())
return;
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;
});
song.setFormats(formats);
}
private void getFormats(Song song)
{
if (!song.getFormats().isEmpty())
return;
try
{
String cmd = Chords.getSettings().getYtdlCommand() + " --skip-download -F " + song.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<>();
List<String> list = output.lines().collect(Collectors.toList());
int i = 0;
while (!list.get(i).contains("Available formats"))
i++;
i++;
for (; i < list.size(); i++)
{
String line = list.get(i);
String[] split = line.split("\\s\\s+", 4);
if (split.length < 4)
continue;
formats.add(new Format(split[0], split[1], split[2], split[3]));
}
song.setFormats(formats);
} catch (Exception ex)
{
Logger.getLogger(Downloader.class.getName()).log(Level.SEVERE, null, ex);
// return List.of();
}
}
private void getInfo(Song song)
{
try
{
String cmd = Chords.getSettings().getYtdlCommand() + " --skip-download --print-json " + song.getUrl().toString();
Process exec = runCommand(cmd, INFO_TIMEOUT);
InputStream input = exec.getInputStream();
JsonReader reader = Json.createReader(input);
JsonObject object = reader.readObject();
if (song.getTitle() == null)
song.setTitle(object.getString("title", null));
if (song.getArtist() == null)
song.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);
}
song.setFormats(formats);
}
} catch (Exception ex)
{
Logger.getLogger(Downloader.class.getName()).log(Level.SEVERE, null, ex);
}
}
private File getDownloadDir() throws IOException
{
if (downloadDir == null || !downloadDir.exists() || !downloadDir.canWrite())
downloadDir = Files.createTempDirectory("chords").toFile();
return downloadDir;
}
private void download(DownloadTask task)
{
Song song = task.request.getSong();
chooseFormats(song);
String formatCodes = "";
final List<Format> formats = song.getFormats();
for (int i = 0; i < 3 && i < song.getFormats().size(); i++)
formatCodes += formats.get(i).getCode() + "/";
try
{
messageHandler.accept(task.request, null);
String cmd = Chords.getSettings().getYtdlCommand()
+ " -x"
+ " -f " + formatCodes + "worstaudio/bestaudio/worst/best"
+ " --audio-format=wav"
+ " --no-playlist"
+ " -o " + getDownloadDir().getAbsolutePath() + "/%(title)s.%(ext)s "
+ song.getUrl().toString();
Process exec = runCommand(cmd, DOWNLOAD_TIMEOUT);
InputStream in = exec.getInputStream();
String output = new String(in.readAllBytes(), Charset.defaultCharset());
String error = new String(exec.getErrorStream().readAllBytes(), Charset.defaultCharset());
System.out.println(output);
Matcher matcher = DESTINATION_PATTERN.matcher(output);
if (matcher.find())
song.setLocation(new File(matcher.group(1)));
else if (exec.exitValue() != 0)
throw new RuntimeException("youtube-dl failed with error " + exec.exitValue() + ", output:\n" + error);
// return true;
if (task.getDestination() != null)
task.getDestination().accept(song);
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();
boolean done = exec.waitFor(timeoutSecs, TimeUnit.SECONDS);
if (!done)
{
exec.destroyForcibly();
throw new RuntimeException("Took too long, giving up.");
}
return exec;
}
public BiConsumer<SongRequest, Exception> getMessageHandler()
{
return messageHandler;
}
public void setMessageHandler(BiConsumer<SongRequest, Exception> messageHandler)
{
this.messageHandler = messageHandler;
}
public List<DownloadTask> getDownloadQueue()
{
return downloadQueue;
}
public static class DownloadTask
{
private final SongRequest request;
private final Consumer<Song> destination;
public DownloadTask(SongRequest request, Consumer<Song> destination)
{
this.request = request;
this.destination = destination;
}
public SongRequest getSong()
{
return request;
}
public Consumer<Song> getDestination()
{
return destination;
}
}
}