125 lines
4.3 KiB
Rust
125 lines
4.3 KiB
Rust
use std::collections::HashMap;
|
|
use std::path::{Path, PathBuf};
|
|
use clap::{Args, Command, Parser, Subcommand};
|
|
use color_eyre::eyre::Result;
|
|
use crossterm::style::{style, Stylize};
|
|
use futures::StreamExt;
|
|
use indicatif::{ProgressBar, ProgressStyle};
|
|
use tokio::time::Instant;
|
|
use crate::models;
|
|
use crate::config::types::ApplicationConfig;
|
|
use crate::constants::{DB_CF_OPTIONS, DB_OPTIONS};
|
|
use crate::crawler::{dlsite, DLSiteCrawler};
|
|
use crate::helpers;
|
|
use crate::helpers::db::RocksDB;
|
|
|
|
#[derive(Parser, Debug)]
|
|
pub(super) struct DLSiteCommand {
|
|
#[command(subcommand)]
|
|
pub(super) subcommand: DLSiteSubCommand,
|
|
}
|
|
|
|
#[derive(Parser, Debug)]
|
|
pub(super) enum DLSiteSubCommand {
|
|
#[command(name = "sync")]
|
|
Sync(DLSiteSyncCommand)
|
|
}
|
|
|
|
#[derive(Parser, Debug)]
|
|
pub(super) struct DLSiteSyncCommand {
|
|
#[clap(long, short, action)]
|
|
missing: bool,
|
|
#[clap(long, short, action)]
|
|
genre: bool,
|
|
}
|
|
|
|
impl DLSiteSubCommand {
|
|
pub async fn handle(&self) -> Result<()> {
|
|
match self {
|
|
Self::Sync(cmd) => cmd.handle().await,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl DLSiteSyncCommand {
|
|
pub async fn handle(&self) -> Result<()> {
|
|
let now = Instant::now();
|
|
let app_conf = ApplicationConfig::get_config()?;
|
|
let mut db = RocksDB::new(DB_OPTIONS.clone(), DB_CF_OPTIONS.clone())?;
|
|
if self.genre {
|
|
Self::sync_genres(&app_conf).await?;
|
|
}
|
|
self.sync_works(&app_conf, &mut db).await?;
|
|
println!("{} Done in {:.2?}", style("Syncing").green(), now.elapsed());
|
|
Ok(())
|
|
}
|
|
|
|
async fn sync_genres(app_conf: &ApplicationConfig) -> Result<()> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn sync_works(&self, app_conf: &ApplicationConfig, db: &mut RocksDB) -> Result<()> {
|
|
let crawler = DLSiteCrawler::new();
|
|
let existing_works = db.get_all_values::<models::DLSiteManiax>()?;
|
|
|
|
let work_list = self.get_work_list(&app_conf, existing_works).await?;
|
|
let rj_nums = work_list.clone().into_keys().collect::<Vec<_>>();
|
|
let mut maniaxes: Vec<models::DLSiteManiax> = Vec::new();
|
|
let mut game_infos = crawler.get_game_infos(rj_nums).await?;
|
|
|
|
let progress = ProgressBar::new(game_infos.len() as u64)
|
|
.with_style(ProgressStyle::default_bar());
|
|
while let Some(info) = game_infos.next().await {
|
|
let mut value: models::DLSiteManiax = info?.into();
|
|
let maniax_folder = work_list.get(&value.rj_num).unwrap().to_owned();
|
|
value.folder_path = maniax_folder;
|
|
maniaxes.push(value);
|
|
progress.inc(1);
|
|
}
|
|
db.set_values(&maniaxes)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn get_work_list(&self, app_conf: &ApplicationConfig, existing_works: Vec<models::DLSiteManiax>) -> Result<HashMap<String, PathBuf>> {
|
|
let existing_nums = existing_works.iter()
|
|
.map(|x| x.rj_num.clone())
|
|
.collect::<Vec<_>>();
|
|
let existing_folders = existing_works.iter()
|
|
.map(|x| x.folder_path.to_str().unwrap().to_string())
|
|
.collect::<Vec<_>>();
|
|
let mut works_list: HashMap<String, PathBuf> = HashMap::new();
|
|
let config_paths = app_conf.path_config.dlsite_paths.iter()
|
|
.map(|path| Path::new(path))
|
|
.collect::<Vec<_>>();
|
|
let dir_paths = helpers::get_all_folders(config_paths).await?;
|
|
for dir_path in dir_paths {
|
|
if !dir_path.is_dir() {
|
|
println!(
|
|
"{} {}",
|
|
style(dir_path.to_str().unwrap()).blue(),
|
|
style("is not a directory").red()
|
|
);
|
|
continue;
|
|
}
|
|
let dir_path_str = dir_path.to_str().unwrap().to_string();
|
|
let dir_name = dir_path
|
|
.file_name().unwrap()
|
|
.to_str().unwrap()
|
|
.to_string();
|
|
if !dlsite::is_valid_rj_number(&dir_name) && !existing_folders.contains(&dir_path_str) {
|
|
println!(
|
|
"{} {}",
|
|
style(dir_path.to_str().unwrap()).blue(),
|
|
style("is not a valid rj number, please add it manually").red()
|
|
);
|
|
continue;
|
|
}
|
|
if self.missing && existing_nums.contains(&dir_name) {
|
|
continue;
|
|
}
|
|
works_list.insert(dir_name, dir_path);
|
|
}
|
|
Ok(works_list)
|
|
}
|
|
} |