mod utils;
use std::{
cmp::Ordering,
path::{Path, PathBuf},
sync::Arc,
};
use anyhow::Result;
use dynasty::utils::{get_dynasty_chapter_identifier, truncate_dynasty_chapter_title};
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use once_cell::sync::Lazy;
use reqwest::{header, Url};
use serde::{Deserialize, Serialize};
use tokio::{
fs::{self, File, OpenOptions},
io::{AsyncWriteExt, BufWriter},
sync::Semaphore,
};
use crate::{
runner::{DefaultRetryConfig, RunnerConfig},
styles::{
ALT_PREFIX_COLOR, BYTES_PROGRESS_COLOR, DOTS_SPINNER_CHARS, DOTS_SPINNER_INTERVAL,
DOWNLOAD_BAR_COLOR, MESSAGE_COLOR, PREFIX_COLOR, PROGRESS_COLOR, SPINNER_COLOR,
},
Cli,
};
use self::utils::{get_content_length_from_response, get_file_extension_from_url};
static DOWNLOADER_MAX_MESSAGE_LENGTH: usize = 53;
static DOWNLOADER_MAIN_PROGRESS_STYLE: Lazy<ProgressStyle> = Lazy::new(|| {
ProgressStyle::with_template(&format!(
"{} {} {} {}",
SPINNER_COLOR.apply_to("{spinner}"),
PREFIX_COLOR.apply_to("{prefix}"),
MESSAGE_COLOR.apply_to("{msg}"),
PROGRESS_COLOR.apply_to("{pos:>2}/{len:<2}")
))
.unwrap()
.tick_chars(DOTS_SPINNER_CHARS)
});
static DOWNLOADER_MAIN_PROGRESS_FINISH_STYLE: Lazy<ProgressStyle> = Lazy::new(|| {
ProgressStyle::with_template(&format!(
"{} {} {} {}",
SPINNER_COLOR.apply_to(">"),
PREFIX_COLOR.apply_to("{prefix}"),
MESSAGE_COLOR.apply_to("{msg}"),
PROGRESS_COLOR.apply_to("{pos:>2}/{len:<2}")
))
.unwrap()
});
static DOWNLOADER_DOWNLOAD_PROGRESS_STYLE: Lazy<ProgressStyle> = Lazy::new(|| {
ProgressStyle::with_template(&format!(
"{} {} |{}| {}",
ALT_PREFIX_COLOR.apply_to("{prefix:>9}"),
BYTES_PROGRESS_COLOR.apply_to("{bytes:>14} {bytes_per_sec:>14}"),
DOWNLOAD_BAR_COLOR.apply_to("{bar:15}"),
PROGRESS_COLOR.apply_to("{percent:>3}%")
))
.unwrap()
.progress_chars("█▇▆▅▄▃▂▁ ")
});
#[derive(Debug, Clone)]
struct Downloader {
#[cfg(feature = "archiver")]
archiver: Option<crate::archiver::Archiver>,
client: reqwest::Client,
semaphore: Arc<Semaphore>,
multi_progress: MultiProgress,
main_progress_bar: ProgressBar,
retry_config: DefaultRetryConfig,
}
impl Downloader {
fn new(client: reqwest::Client, config: &RunnerConfig) -> Self {
let multi_progress = MultiProgress::new();
let main_progress_bar = multi_progress
.add(ProgressBar::new(0).with_style(DOWNLOADER_MAIN_PROGRESS_STYLE.clone()));
main_progress_bar.enable_steady_tick(DOTS_SPINNER_INTERVAL);
Downloader {
client,
multi_progress,
main_progress_bar,
semaphore: config.semaphore.clone(),
retry_config: config.retry_config,
#[cfg(feature = "archiver")]
archiver: config.archiver,
}
}
fn add_progress_bar(&self) -> ProgressBar {
self.multi_progress.insert_before(
&self.main_progress_bar,
ProgressBar::new(1).with_style(DOWNLOADER_DOWNLOAD_PROGRESS_STYLE.clone()),
)
}
async fn download_chapter(&self, chapter: DownloadableChapter, path: PathBuf) -> Result<()> {
let chapter_title = chapter.title.clone();
let update_progress_bar_prefix = |prefix: &'static str| {
self.main_progress_bar.set_prefix(prefix);
self.main_progress_bar
.set_message(truncate_dynasty_chapter_title(
&chapter_title,
DOWNLOADER_MAX_MESSAGE_LENGTH - prefix.len(),
));
};
let length = chapter.pages.len();
self.main_progress_bar.set_length(length as u64);
self.main_progress_bar.set_position(0);
update_progress_bar_prefix("downloading");
let mut handles = Vec::with_capacity(length);
for (index, url) in chapter.pages.into_iter().enumerate() {
let url = Url::parse(&url)?;
let path = path.clone();
let self_clone = self.clone();
let permit = self.semaphore.clone().acquire_owned().await?;
handles.push(tokio::spawn(async move {
let page_path = self_clone.download_url(url, index + 1, path).await?;
drop(permit);
self_clone.main_progress_bar.inc(1);
anyhow::Ok(page_path)
}))
}
let mut pages = Vec::with_capacity(length);
for handle in handles {
pages.push(handle.await??);
}
update_progress_bar_prefix("downloaded");
#[cfg(feature = "archiver")]
{
use crate::archiver::ArchivableChapter;
if let Some(archiver) = self.archiver {
update_progress_bar_prefix("archiving");
let archivable = ArchivableChapter {
pages,
path,
title: chapter.title,
};
archiver.archive(&archivable, &self.main_progress_bar)?;
update_progress_bar_prefix("archived");
}
}
self.main_progress_bar.set_position(length as u64);
self.main_progress_bar
.set_style(DOWNLOADER_MAIN_PROGRESS_FINISH_STYLE.clone());
self.main_progress_bar.finish();
Ok(())
}
async fn download_url(&self, url: Url, index: usize, path: PathBuf) -> Result<PathBuf> {
let extension = get_file_extension_from_url(&url).unwrap_or_else(|| "png".to_string());
let filename = format!("{:0>3}.{}", index, extension);
let page_path = path.join(&filename);
if page_path.exists() {
return Ok(page_path);
}
let progress_bar = self.add_progress_bar();
progress_bar.set_prefix(filename.clone());
let part_filename = format!("{}.part", filename);
let part_page_path = page_path.with_file_name(&part_filename);
let (file, download_range) = if part_page_path.exists() {
let existing_part_file = OpenOptions::new()
.append(true)
.open(&part_page_path)
.await?;
let downloaded = existing_part_file.metadata().await?.len();
let head_response = tryhard::retry_fn(|| async {
self.client
.head(url.as_str())
.send()
.await
.and_then(|response| response.error_for_status())
})
.with_config(self.retry_config)
.await?;
let content_length = get_content_length_from_response(&head_response);
if content_length == 0 || downloaded > content_length {
(File::create(&part_page_path).await?, None)
} else {
match downloaded.cmp(&content_length) {
Ordering::Equal => {
fs::rename(&part_page_path, &page_path).await?;
progress_bar.finish_and_clear();
return Ok(page_path);
}
Ordering::Less => {
progress_bar.set_length(content_length);
progress_bar.set_position(downloaded);
progress_bar.reset_eta();
(existing_part_file, Some(downloaded..content_length))
}
Ordering::Greater => unreachable!(),
}
}
} else {
(File::create(&part_page_path).await?, None)
};
let request = {
let builder = self.client.get(url.as_str());
if let Some(range) = download_range {
builder.header(
header::RANGE,
format!("bytes={}-{}", range.start, range.end),
)
} else {
builder
}
};
let mut response = tryhard::retry_fn(|| {
let request = request.try_clone().unwrap();
async {
request
.send()
.await
.and_then(|response| response.error_for_status())
}
})
.with_config(self.retry_config)
.await?;
if progress_bar.length() == Some(1) {
let content_length = get_content_length_from_response(&response);
progress_bar.set_length(content_length);
}
let mut writer = BufWriter::new(file);
while let Some(bytes) = response.chunk().await? {
writer.write_all(&bytes).await?;
progress_bar.inc(bytes.len() as u64);
}
writer.shutdown().await?;
fs::rename(&part_page_path, &page_path).await?;
progress_bar.finish_and_clear();
Ok(page_path)
}
}
#[derive(Debug)]
pub struct Downloadable {
pub title: Option<String>,
pub chapters: Vec<DownloadableChapter>,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct DownloadableChapter {
pub title: String,
pub pages: Vec<String>,
}
pub async fn download(
downloadable: Downloadable,
client: reqwest::Client,
cli: &Cli,
config: &RunnerConfig,
) -> Result<()> {
for chapter in downloadable.chapters {
let chapter_path = {
let base = cli.directory.as_deref().unwrap_or_else(|| Path::new("."));
if cli.flatten_directory || downloadable.title.is_none() {
base.join(&chapter.title)
} else {
let path = get_dynasty_chapter_identifier(&chapter.title).unwrap_or(&chapter.title);
base.join(downloadable.title.as_ref().unwrap()).join(path)
}
};
if !chapter_path.exists() {
fs::create_dir_all(&chapter_path).await?;
}
Downloader::new(client.clone(), config)
.download_chapter(chapter, chapter_path)
.await?;
}
Ok(())
}