package io.lumify.themoviedb.download;
import io.lumify.core.cmdline.CommandLineBase;
import io.lumify.core.util.JSONUtil;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.OptionBuilder;
import org.apache.commons.cli.Options;
import org.apache.commons.io.FileUtils;
import org.json.JSONArray;
import org.json.JSONObject;
import java.io.File;
import java.io.IOException;
import java.text.ParseException;
import java.util.LinkedList;
import java.util.Queue;
import java.util.Random;
public class TheMovieDbDownload extends CommandLineBase {
private static final String CMD_OPT_API_KEY = "apikey";
private static final String CMD_OPT_CACHE_DIRECTORY = "cachedir";
private static final String CMD_OPT_DOWNLOAD_PERSON = "personid";
private static final String CMD_OPT_DOWNLOAD_MOVIE = "movieid";
private static final String CMD_OPT_DOWNLOAD_PRODUCTION_COMPANY = "productioncompanyid";
public static final String DIR_MOVIES = "movies";
public static final String DIR_PERSONS = "persons";
public static final String DIR_IMAGES = "images";
public static final String DIR_PRODUCTION_COMPANIES = "productionCompanies";
private final Queue<WorkItem> workQueue = new LinkedList<WorkItem>();
private TheMovieDb theMovieDb;
private File cacheDir;
private File moviesDir;
private File personsDir;
private File imagesDir;
private File productionCompanyDir;
public static void main(String[] args) throws Exception {
int res = new TheMovieDbDownload().run(args);
if (res != 0) {
System.exit(res);
}
}
@Override
protected Options getOptions() {
Options opts = super.getOptions();
opts.addOption(
OptionBuilder
.withLongOpt(CMD_OPT_API_KEY)
.withDescription("TheMovieDb API Key")
.hasArg()
.isRequired()
.create()
);
opts.addOption(
OptionBuilder
.withLongOpt(CMD_OPT_CACHE_DIRECTORY)
.withDescription("Directory to cache json documents in")
.hasArg()
.isRequired()
.create()
);
opts.addOption(
OptionBuilder
.withLongOpt(CMD_OPT_DOWNLOAD_PERSON)
.withDescription("Initial person id to download")
.hasArg()
.create()
);
opts.addOption(
OptionBuilder
.withLongOpt(CMD_OPT_DOWNLOAD_MOVIE)
.withDescription("Initial movie id to download")
.hasArg()
.create()
);
opts.addOption(
OptionBuilder
.withLongOpt(CMD_OPT_DOWNLOAD_PRODUCTION_COMPANY)
.withDescription("Initial production company id to download")
.hasArg()
.create()
);
return opts;
}
@Override
protected int run(CommandLine cmd) throws Exception {
String apiKey = cmd.getOptionValue(CMD_OPT_API_KEY);
theMovieDb = new TheMovieDb(apiKey);
String cacheDirString = cmd.getOptionValue(CMD_OPT_CACHE_DIRECTORY);
this.cacheDir = new File(cacheDirString);
cacheDir.mkdirs();
this.moviesDir = new File(cacheDir, DIR_MOVIES);
moviesDir.mkdirs();
this.personsDir = new File(cacheDir, DIR_PERSONS);
personsDir.mkdirs();
this.imagesDir = new File(cacheDir, DIR_IMAGES);
imagesDir.mkdirs();
this.productionCompanyDir = new File(cacheDir, DIR_PRODUCTION_COMPANIES);
productionCompanyDir.mkdirs();
String[] personIds = cmd.getOptionValues(CMD_OPT_DOWNLOAD_PERSON);
if (personIds != null) {
for (String personIdString : personIds) {
int personId = Integer.parseInt(personIdString);
queuePersonDownload(personId);
}
}
String[] movieIds = cmd.getOptionValues(CMD_OPT_DOWNLOAD_MOVIE);
if (movieIds != null) {
for (String movieIdString : movieIds) {
int movieId = Integer.parseInt(movieIdString);
queueMovieDownload(movieId);
}
}
String[] productionCompanyIds = cmd.getOptionValues(CMD_OPT_DOWNLOAD_PRODUCTION_COMPANY);
if (productionCompanyIds != null) {
for (String productionCompanyIdString : productionCompanyIds) {
int productionCompanyId = Integer.parseInt(productionCompanyIdString);
queueProductionCompanyDownload(productionCompanyId);
}
}
loadCacheDir();
download();
return 0;
}
private void download() {
Random random = new Random();
while (true) {
LOGGER.debug("%d still in queue", workQueue.size());
WorkItem workItem = workQueue.poll();
if (workItem == null) {
break;
}
try {
if (workItem.process(this)) {
Thread.sleep(1000 + random.nextInt(500));
}
} catch (Exception e) {
LOGGER.error("Could not process work item: %s", workItem.toString(), e);
}
}
}
private void loadCacheDir() {
LOGGER.debug("Loading: %s", getCacheDir().getAbsolutePath());
loadProductionCompaniesCacheDir(getProductionCompanyDir());
loadPersonsCacheDir(getPersonsDir());
loadMoviesCacheDir(getMoviesDir());
}
private void loadMoviesCacheDir(File moviesDir) {
File[] files = moviesDir.listFiles();
int i = 1;
for (File movieFile : files) {
LOGGER.debug("loading movie %d/%d", i, files.length);
loadMovieFile(movieFile);
i++;
}
}
private void loadMovieFile(File movieFile) {
try {
if (!movieFile.getName().endsWith(".json")) {
return;
}
LOGGER.debug("Loading Movie File: %s", movieFile);
String data = FileUtils.readFileToString(movieFile);
JSONObject movieJson = JSONUtil.parse(data);
loadMovieJson(movieJson);
} catch (Exception e) {
LOGGER.error("Could not read movie %s", movieFile.getAbsolutePath(), e);
}
}
private void loadMovieJson(JSONObject movieJson) throws IOException, ParseException {
JSONObject credits = movieJson.getJSONObject("credits");
JSONArray cast = credits.getJSONArray("cast");
for (int i = 0; i < cast.length(); i++) {
JSONObject castJson = cast.getJSONObject(i);
int personId = castJson.getInt("id");
queuePersonDownload(personId);
}
JSONArray productionCompanies = movieJson.optJSONArray("production_companies");
if (productionCompanies != null) {
for (int i = 0; i < productionCompanies.length(); i++) {
JSONObject productionCompany = productionCompanies.getJSONObject(i);
int productionCompanyId = productionCompany.getInt("id");
queueProductionCompanyDownload(productionCompanyId);
}
}
getGraph().flush();
String posterImagePath = movieJson.optString("poster_path");
if (posterImagePath != null && posterImagePath.length() > 0) {
if (!hasImageInCache(posterImagePath)) {
queueImageDownload(posterImagePath);
}
}
}
private void loadPersonsCacheDir(File personsDir) {
File[] files = personsDir.listFiles();
int i = 1;
for (File personFile : files) {
LOGGER.debug("Loading person %d/%d", i, files.length);
loadPersonFile(personFile);
i++;
}
}
private void loadPersonFile(File personFile) {
try {
if (!personFile.getName().endsWith(".json")) {
return;
}
LOGGER.debug("Loading Person File: %s", personFile);
String data = FileUtils.readFileToString(personFile);
JSONObject personJson = JSONUtil.parse(data);
loadPersonJson(personJson);
} catch (Exception e) {
LOGGER.error("Could not read person %s", personFile.getAbsolutePath(), e);
}
}
private void loadPersonJson(JSONObject personJson) throws ParseException, IOException {
JSONObject combinedCredits = personJson.getJSONObject("combined_credits");
JSONArray cast = combinedCredits.getJSONArray("cast");
for (int i = 0; i < cast.length(); i++) {
JSONObject movieJson = cast.getJSONObject(i);
String mediaType = movieJson.getString("media_type");
if (!mediaType.equals("movie")) {
continue;
}
int movieId = movieJson.getInt("id");
queueMovieDownload(movieId);
}
getGraph().flush();
String profileImage = personJson.optString("profile_path");
if (profileImage != null && profileImage.length() > 0) {
if (!hasImageInCache(profileImage)) {
queueImageDownload(profileImage);
}
}
}
private void loadProductionCompaniesCacheDir(File productionCompanyDir) {
for (File productionCompanyFile : productionCompanyDir.listFiles()) {
loadProductionCompanyFile(productionCompanyFile);
}
}
private void loadProductionCompanyFile(File productionCompanyFile) {
try {
if (!productionCompanyFile.getName().endsWith(".json")) {
return;
}
LOGGER.debug("Loading production company File: %s", productionCompanyFile);
String data = FileUtils.readFileToString(productionCompanyFile);
JSONObject productionCompanyJson = JSONUtil.parse(data);
loadProductionCompanyJson(productionCompanyJson);
} catch (Exception e) {
LOGGER.error("Could not read production company %s", productionCompanyFile.getAbsolutePath(), e);
}
}
private void loadProductionCompanyJson(JSONObject productionCompanyJson) throws IOException {
int productionCompanyId = productionCompanyJson.getInt("id");
LOGGER.debug("Loading production company: %d", productionCompanyId);
String logoImagePath = productionCompanyJson.optString("logo_path");
if (logoImagePath != null && logoImagePath.length() > 0) {
if (!hasImageInCache(logoImagePath)) {
queueImageDownload(logoImagePath);
}
}
}
public File getCacheDir() {
return cacheDir;
}
public File getMoviesDir() {
return moviesDir;
}
public File getPersonsDir() {
return personsDir;
}
public File getProductionCompanyDir() {
return productionCompanyDir;
}
public File getImagesDir() {
return imagesDir;
}
public boolean hasPersonInCache(int personId) {
return new File(getPersonsDir(), personId + ".json").exists();
}
public boolean hasProductionCompanyInCache(int productionCompanyId) {
return new File(getProductionCompanyDir(), productionCompanyId + ".json").exists();
}
public boolean hasMovieInCache(int movieId) {
return new File(getMoviesDir(), movieId + ".json").exists();
}
public boolean hasImageInCache(String imagePath) {
return new File(getImagesDir(), imagePath).exists();
}
private void queueImageDownload(String imagePath) {
if (hasImageInCache(imagePath)) {
return;
}
workQueue.add(new ImageDownloadWorkItem(imagePath));
}
private void queueMovieDownload(int movieId) {
if (hasMovieInCache(movieId)) {
return;
}
workQueue.add(new MovieDownloadWorkItem(movieId));
}
private void queuePersonDownload(int personId) {
if (hasPersonInCache(personId)) {
return;
}
workQueue.add(new PersonDownloadWorkItem(personId));
}
private void queueProductionCompanyDownload(int productionCompanyId) {
if (hasProductionCompanyInCache(productionCompanyId)) {
return;
}
workQueue.add(new ProductionCompanyDownloadWorkItem(productionCompanyId));
}
public TheMovieDb getTheMovieDb() {
return theMovieDb;
}
public void writeMovie(int movieId, JSONObject movieJson) throws IOException, ParseException {
File movieFile = new File(getMoviesDir(), movieId + ".json");
FileUtils.write(movieFile, movieJson.toString(2));
loadMovieJson(movieJson);
}
public void writePerson(int personId, JSONObject personJson) throws IOException, ParseException {
File personFile = new File(getPersonsDir(), personId + ".json");
FileUtils.write(personFile, personJson.toString(2));
loadPersonJson(personJson);
}
public void writeProductionCompany(int productionCompanyId, JSONObject productionCompanyJson) throws IOException {
File personFile = new File(getProductionCompanyDir(), productionCompanyId + ".json");
FileUtils.write(personFile, productionCompanyJson.toString(2));
loadProductionCompanyJson(productionCompanyJson);
}
public void writeImage(String imagePath, byte[] imageData) throws IOException {
File imageFile = new File(getImagesDir(), imagePath);
imageFile.getParentFile().mkdirs();
FileUtils.writeByteArrayToFile(imageFile, imageData);
}
}