1use std::collections::BTreeSet;
2
3use anyhow::Result;
4use fs_err::File;
5use futures_channel::mpsc;
6
7use abstio::{DataPacks, Manifest, MapName};
8use abstutil::prettyprint_bytes;
9use widgetry::tools::{ChooseSomething, FutureLoader, PopupMsg};
10use widgetry::{EventCtx, Key, Transition};
11
12use crate::AppLike;
13
14pub fn prompt_to_download_missing_data<A: AppLike + 'static>(
15 ctx: &mut EventCtx,
16 map_name: MapName,
17 on_load: Box<dyn FnOnce(&mut EventCtx, &mut A) -> Transition<A>>,
18) -> Transition<A> {
19 let manifest = files_to_download(&map_name);
20 let bytes = manifest
21 .entries
22 .iter()
23 .map(|(_, e)| e.compressed_size_bytes)
24 .sum();
25
26 Transition::Push(ChooseSomething::new_state(
27 ctx,
28 format!(
29 "Missing data. Download {} for {}?",
30 prettyprint_bytes(bytes),
31 map_name.describe()
32 ),
33 vec![
34 widgetry::Choice::string("Yes, download"),
35 widgetry::Choice::string("Never mind").key(Key::Escape),
36 ],
37 Box::new(move |resp, ctx, _| {
38 if resp == "Never mind" {
39 return Transition::Pop;
40 }
41
42 let (outer_progress_tx, outer_progress_rx) = futures_channel::mpsc::channel(1000);
43 let (inner_progress_tx, inner_progress_rx) = futures_channel::mpsc::channel(1000);
44 Transition::Replace(FutureLoader::<A, Result<()>>::new_state(
45 ctx,
46 Box::pin(async {
47 let result =
48 download_files(manifest, outer_progress_tx, inner_progress_tx).await;
49 let wrap: Box<dyn Send + FnOnce(&A) -> Result<()>> =
50 Box::new(move |_: &A| result);
51 Ok(wrap)
52 }),
53 outer_progress_rx,
54 inner_progress_rx,
55 "Downloading missing files",
56 Box::new(|ctx, app, maybe_result| {
57 let error_msg = match maybe_result {
58 Ok(Ok(())) => None,
59 Ok(Err(err)) => Some(err.to_string()),
60 Err(err) => Some(format!("Something went very wrong: {}", err)),
61 };
62 if let Some(err) = error_msg {
63 Transition::Replace(PopupMsg::new_state(ctx, "Download failed", vec![err]))
64 } else {
65 on_load(ctx, app)
66 }
67 }),
68 ))
69 }),
70 ))
71}
72
73fn files_to_download(map: &MapName) -> Manifest {
74 let mut data_packs = DataPacks {
75 runtime: BTreeSet::new(),
76 input: BTreeSet::new(),
77 };
78 data_packs.runtime.insert(map.to_data_pack_name());
79 let mut manifest = Manifest::load().filter(data_packs);
80 manifest
82 .entries
83 .retain(|path, _| !abstio::file_exists(&abstio::path(path.strip_prefix("data/").unwrap())));
84
85 manifest.entries.retain(|path, _| {
88 let parts = path.split('/').collect::<Vec<_>>();
90 parts[4] == "city.bin"
91 || (parts[4] == "maps" && parts[5] == format!("{}.bin", map.map))
92 || (parts.len() >= 6 && parts[5] == map.map)
93 });
94
95 manifest
96}
97
98async fn download_files(
99 manifest: Manifest,
100 mut outer_progress: mpsc::Sender<String>,
101 mut inner_progress: mpsc::Sender<String>,
102) -> Result<()> {
103 let num_files = manifest.entries.len();
104 let mut messages = Vec::new();
105 let mut files_so_far = 0;
106
107 for (path, entry) in manifest.entries {
108 files_so_far += 1;
109 let local_path = abstio::path(path.strip_prefix("data/").unwrap());
110 let url = format!(
111 "https://play.abstreet.org/{}/{}.gz",
112 crate::tools::version(),
113 path
114 );
115 if let Err(err) = outer_progress.try_send(format!(
116 "Downloading file {}/{}: {} ({})",
117 files_so_far,
118 num_files,
119 url,
120 prettyprint_bytes(entry.compressed_size_bytes)
121 )) {
122 warn!("Couldn't send progress: {}", err);
123 }
124
125 match abstio::download_bytes(&url, None, &mut inner_progress)
126 .await
127 .and_then(|bytes| {
128 info!("Decompressing {}", path);
131 fs_err::create_dir_all(std::path::Path::new(&local_path).parent().unwrap())
132 .unwrap();
133 let mut out = File::create(&local_path).unwrap();
134 let mut decoder = flate2::read::GzDecoder::new(&bytes[..]);
135 std::io::copy(&mut decoder, &mut out).map_err(|err| err.into())
136 }) {
137 Ok(_) => {}
138 Err(err) => {
139 let msg = format!("Problem with {}: {}", url, err);
140 error!("{}", msg);
141 messages.push(msg);
142 }
143 }
144 }
145 if !messages.is_empty() {
146 bail!("{} errors: {}", messages.len(), messages.join(", "));
147 }
148 Ok(())
149}