| 1 | use crate::{ |
| 2 | audio_dir_path, command::player_prepare::get_mode_from_filename, models::player::ResultFile, players::{BASS_FILE, DRUMS_FILE, OTHER_FILE, VOCALS_FILE} |
| 3 | }; |
| 4 | use anyhow::anyhow; |
| 5 | use futures::future; |
| 6 | use log::{debug, error, info, warn}; |
| 7 | use reqwest::{Client, Error, Response}; |
| 8 | use std::{ |
| 9 | fs::{create_dir_all, File, OpenOptions}, |
| 10 | io::Write, |
| 11 | path::{Path, PathBuf}, |
| 12 | }; |
| 13 | use super::player_set_volume::Mode; |
| 14 | |
| 15 | pub async fn save_result_files(audio_file_id: &String, result_files: &[ResultFile]) { |
| 16 | // Check if we have ~/.sfai/audio/{provider_id} folder and create it if not |
| 17 | let audio_id_path = match get_provider_id_dir_path(audio_file_id).await { |
| 18 | Ok(path) => path, |
| 19 | Err(e) => { |
| 20 | error!("Failed to get audio directory: {:?}", e); |
| 21 | return; |
| 22 | } |
| 23 | }; |
| 24 | |
| 25 | // Write length and onset to result file's filename.json |
| 26 | let x = result_files |
| 27 | .iter() |
| 28 | .map(|result_file| { |
| 29 | persist_result_file(result_file.clone()); |
| 30 | let filename = format!("{}.json", result_file.filename); |
| 31 | let file = get_result_file(&audio_id_path, &filename, true); |
| 32 | async move { |
| 33 | match file { |
| 34 | Ok(mut file) => { |
| 35 | let _ = |
| 36 | file.write_all(serde_json::to_string(&result_file).unwrap().as_bytes()); |
| 37 | } |
| 38 | Err(e) => { |
| 39 | error!("Failed to get file: {:?}", e); |
| 40 | } |
| 41 | } |
| 42 | } |
| 43 | }) |
| 44 | .collect::<Vec<_>>(); |
| 45 | |
| 46 | let _ = future::join_all(x).await; |
| 47 | info!("Saved audio analysis."); |
| 48 | } |
| 49 | |
| 50 | fn persist_result_file(result_file: ResultFile) { |
| 51 | debug!("Persisting result file: {}", result_file); |
| 52 | let mode = get_mode_from_filename(&result_file.filename); |
| 53 | match mode { |
| 54 | Mode::Vocals => { |
| 55 | let _ = unsafe { VOCALS_FILE.set(result_file) }; |
| 56 | } |
| 57 | Mode::Bass => { |
| 58 | let _ = unsafe { BASS_FILE.set(result_file) }; |
| 59 | } |
| 60 | Mode::Other => { |
| 61 | let _ = unsafe { OTHER_FILE.set(result_file) }; |
| 62 | } |
| 63 | Mode::Drums => { |
| 64 | let _ = unsafe { DRUMS_FILE.set(result_file) }; |
| 65 | } |
| 66 | Mode::Unknown => { |
| 67 | error!("Unknown mode."); |
| 68 | } |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | pub async fn download_result_files(provider_id: String, result_files: Vec<ResultFile>) { |
| 73 | // Check if we have ~/.sfai/audio/{provider_id} folder and create it if not |
| 74 | let provider_id_path = match get_provider_id_dir_path(&provider_id).await { |
| 75 | Ok(path) => path, |
| 76 | Err(e) => { |
| 77 | error!("Failed to get audio directory: {:?}", e); |
| 78 | return; |
| 79 | } |
| 80 | }; |
| 81 | |
| 82 | batch_download(result_files, provider_id_path).await; |
| 83 | } |
| 84 | |
| 85 | async fn batch_download(result_files: Vec<ResultFile>, provider_id_path: PathBuf) { |
| 86 | // Download files in parallel and store them in the ~/.sfai/audio/{provider_id} folder |
| 87 | info!("Batch downloading {:?} result files.", result_files.len()); |
| 88 | |
| 89 | let download_file = |result_file: ResultFile| { |
| 90 | info!("Downloading file: {:?}", result_file.source_file); |
| 91 | let task = async { |
| 92 | let response = Client::new().get(result_file.source_file).send().await; |
| 93 | |
| 94 | match response { |
| 95 | Ok(response) => { |
| 96 | info!("Got response: {:?}", response); |
| 97 | Ok(response) |
| 98 | } |
| 99 | Err(e) => { |
| 100 | error!("Failed to get response: {:?}", e); |
| 101 | Err(e) |
| 102 | } |
| 103 | } |
| 104 | }; |
| 105 | |
| 106 | // Create file in ~/.sfai/audio/{provider_id} folder |
| 107 | let file = get_result_file(&provider_id_path, &result_file.filename, true); |
| 108 | |
| 109 | (task, file) |
| 110 | }; |
| 111 | |
| 112 | let (tasks, files): (Vec<_>, Vec<_>) = result_files |
| 113 | .into_iter() |
| 114 | .filter(|x| { |
| 115 | debug!("Length: {:?}", x.length); |
| 116 | debug!("Onset: {:?}", x.onset); |
| 117 | should_download_file(&provider_id_path, x.filename.clone()) |
| 118 | }) |
| 119 | .map(download_file) |
| 120 | .unzip(); |
| 121 | |
| 122 | let results = future::join_all(tasks).await; |
| 123 | |
| 124 | let store_result = |x: (Result<Response, Error>, Result<File, std::io::Error>)| async { |
| 125 | //info!("Storing result: {:?}, {:?}", x.0, x.1); |
| 126 | |
| 127 | // make sure both are Oks |
| 128 | let (result, file) = match (x.0, x.1) { |
| 129 | (Ok(result), Ok(file)) => { |
| 130 | info!("Got result and file."); |
| 131 | info!("Result: {:?}", result.url()); |
| 132 | info!("File: {:?}", file); |
| 133 | (result, file) |
| 134 | } |
| 135 | _ => { |
| 136 | error!("Failed to get result and file."); |
| 137 | return; |
| 138 | } |
| 139 | }; |
| 140 | // Write bytes to file |
| 141 | let _ = write_bytes_to_file(file, result).await; |
| 142 | }; |
| 143 | |
| 144 | // Map result and file to store_result |
| 145 | let store_results = results |
| 146 | .into_iter() |
| 147 | .zip(files) |
| 148 | .map(|x| async { |
| 149 | debug!("Storing result: {:?}, {:?}", x.0, x.1); |
| 150 | let _ = store_result(x).await; |
| 151 | }) |
| 152 | .collect::<Vec<_>>(); |
| 153 | |
| 154 | let _ = future::join_all(store_results).await; |
| 155 | |
| 156 | info!("Batch download done."); |
| 157 | } |
| 158 | |
| 159 | // Check if we should download the file. |
| 160 | // If we have the ~/.sfai/audio/{provider_id} folder, then we should not download the result files again. |
| 161 | pub async fn get_provider_id_dir_path(provider_id: &String) -> anyhow::Result<PathBuf> { |
| 162 | if let Some(audio_dir_path) = audio_dir_path() { |
| 163 | let provider_dir_path = audio_dir_path.join(provider_id.clone()); |
| 164 | if provider_dir_path.exists() && provider_dir_path.is_dir() { |
| 165 | return Ok(provider_dir_path); |
| 166 | } |
| 167 | // CREATE DIRECTORY |
| 168 | let _ = create_dir_all(audio_dir_path.join(provider_id)); |
| 169 | Ok(provider_dir_path) |
| 170 | } else { |
| 171 | error!("Failed to get audio directory."); |
| 172 | Err(anyhow!("Failed to get audio directory.")) |
| 173 | } |
| 174 | } |
| 175 | |
| 176 | fn get_result_file( |
| 177 | provider_id_path: &Path, |
| 178 | filename: &String, |
| 179 | should_create_file: bool, |
| 180 | ) -> Result<File, std::io::Error> { |
| 181 | OpenOptions::new() |
| 182 | .create(should_create_file) |
| 183 | .write(true) |
| 184 | .read(true) |
| 185 | .open(provider_id_path.join(filename)) |
| 186 | } |
| 187 | |
| 188 | fn should_download_file(provider_id_path: &Path, filename: String) -> bool { |
| 189 | match get_result_file(provider_id_path, &filename, false) { |
| 190 | Ok(file) => { |
| 191 | warn!( |
| 192 | "File already exists: {:?}", |
| 193 | &provider_id_path.join(filename) |
| 194 | ); |
| 195 | // Check if file is not empty |
| 196 | let metadata = file.metadata(); |
| 197 | match metadata { |
| 198 | Ok(metadata) => { |
| 199 | if metadata.len() > 0 { |
| 200 | warn!("File is not empty. We should not download it."); |
| 201 | return false; |
| 202 | } |
| 203 | debug!("File is empty. We should download it."); |
| 204 | true |
| 205 | } |
| 206 | Err(e) => { |
| 207 | error!("Failed to get metadata: {:?}", e); |
| 208 | true |
| 209 | } |
| 210 | } |
| 211 | } |
| 212 | Err(e) => { |
| 213 | debug!( |
| 214 | "Could not find a result file {:?}: {:?}", |
| 215 | &provider_id_path.join(filename), |
| 216 | e.to_string() |
| 217 | ); |
| 218 | true |
| 219 | } |
| 220 | } |
| 221 | } |
| 222 | |
| 223 | // Write reqwest::Response bytes to file. |
| 224 | async fn write_bytes_to_file(mut file: File, response: Response) -> Result<(), std::io::Error> { |
| 225 | info!("Writing bytes to: {:?}", file); |
| 226 | |
| 227 | match response.bytes().await { |
| 228 | Ok(bytes) => match file.write_all(&bytes) { |
| 229 | Ok(_) => { |
| 230 | info!("Wrote bytes to file."); |
| 231 | Ok(()) |
| 232 | } |
| 233 | Err(e) => { |
| 234 | error!("Failed to write bytes to file: {:?}", e); |
| 235 | Err(e) |
| 236 | } |
| 237 | }, |
| 238 | Err(e) => { |
| 239 | error!("Failed to get bytes: {:?}", e); |
| 240 | Err(std::io::Error::new( |
| 241 | std::io::ErrorKind::Other, |
| 242 | "Failed to get bytes.", |
| 243 | )) |
| 244 | } |
| 245 | } |
| 246 | } |