Files
sus-manager/src/cli/sync.rs
2025-10-27 20:37:42 +08:00

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)
}
}