1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use clap::Parser;
use futures::future::join_all;
use indicatif::ProgressBar;
use indicatif::{MultiProgress, ProgressStyle};
use std::fs::create_dir_all;
use std::sync::Arc;
use std::time::Instant;
use tokio::task::JoinHandle;

use crate::download::*;
use crate::fetch::*;

#[derive(Parser, Debug, Clone)]
pub struct Args {
    //https://stackoverflow.com/questions/74936109/how-to-use-clap-to-take-create-a-vector
    /// Space-delimited list of image classes
    #[arg(short, long, num_args = 1.., value_delimiter = ' ', value_parser)]
    pub topics: Vec<String>,

    /// number of images to download per-class
    #[arg(short, long, default_value_t = 20)]
    pub nsamples: usize,

    /// name of directory to save to
    #[arg(short, long, default_value = "images")]
    pub dir: String,
}

struct TopicHandler<'a> {
    topic: &'a str,
    download_manager: DLManager,
}

impl<'a> TopicHandler<'a> {
    pub async fn download_topic(&mut self, dir: String) -> u8 {
        let nsamples = self.download_manager.target_size;

        //download path = base_dir + topic
        let dir = format!("{}/{}", dir, self.topic);

        let batch = Fetcher::query_api(self.topic, nsamples).await;
        let total = self.download_manager.download_batch(batch, &dir).await;
        total
    }
}

pub struct TinyDataClient {
    args: Args,
    multi_progress_bar: MultiProgress,
}

impl TinyDataClient {
    pub fn new(args: Args) -> Self {
        Self {
            args,
            multi_progress_bar: MultiProgress::new(),
        }
    }

    pub async fn run(&mut self) {
        let args = self.args.clone();

        //create image directories
        for topic in &args.topics {
            let dir = format!("{}/{}", args.dir, topic);
            create_dir_all(&dir).expect("failed to create directory");
        }

        let nsamples = args.nsamples;
        // need thread-safe string for async
        let dir = Arc::new(args.dir);

        let futures: Vec<JoinHandle<u8>> = args
            .topics
            .into_iter()
            .map(|topic| {
                let dir = Arc::clone(&dir);
                let mut pb = self
                    .multi_progress_bar
                    .add(ProgressBar::new(nsamples as u64));

                //style
                stylize_pb(&mut pb, &topic);

                tokio::spawn(async move {
                    let download_manager = DLManager::new(nsamples, pb);

                    let mut topic_handler = TopicHandler {
                        topic: &topic,
                        download_manager,
                    };
                    topic_handler.download_topic(dir.to_string()).await
                })
            })
            .collect();

        //time execution
        let now = Instant::now();
        let total = join_all(futures).await;
        let elapsed = now.elapsed();

        let total: u8 = total.into_iter().map(|res| res.unwrap()).sum();

        println!(
            "{}/{} files saved successfully to `./{}` in {}s 📦",
            total,
            self.args.nsamples * self.args.topics.len(),
            self.args.dir,
            elapsed.as_secs(),
        );
    }
}

fn stylize_pb(pb: &mut ProgressBar, name: &str) {
    let default = "{msg} {spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})";
    pb.set_style(
        ProgressStyle::with_template(default)
            .unwrap()
            .progress_chars("#>-"),
    );
    pb.set_message(String::from(name));
}