1use std::path::{Path, PathBuf};
2
3use base64::{engine::general_purpose, Engine as _};
4use rand::{distributions::Alphanumeric, Rng};
5use reqwest::Url;
6
7use crate::error::OpenAIError;
8
9fn create_paths<P: AsRef<Path>>(url: &Url, base_dir: P) -> (PathBuf, PathBuf) {
10 let mut dir = PathBuf::from(base_dir.as_ref());
11 let mut path = dir.clone();
12 let segments = url.path_segments().map(|c| c.collect::<Vec<_>>());
13 if let Some(segments) = segments {
14 for (idx, segment) in segments.iter().enumerate() {
15 if idx != segments.len() - 1 {
16 dir.push(segment);
17 }
18 path.push(segment);
19 }
20 }
21
22 (dir, path)
23}
24
25pub(crate) async fn download_url<P: AsRef<Path>>(
26 url: &str,
27 dir: P,
28) -> Result<PathBuf, OpenAIError> {
29 let parsed_url = Url::parse(url).map_err(|e| OpenAIError::FileSaveError(e.to_string()))?;
30 let response = reqwest::get(url)
31 .await
32 .map_err(|e| OpenAIError::FileSaveError(e.to_string()))?;
33
34 if !response.status().is_success() {
35 return Err(OpenAIError::FileSaveError(format!(
36 "couldn't download file, status: {}, url: {url}",
37 response.status()
38 )));
39 }
40
41 let (dir, file_path) = create_paths(&parsed_url, dir);
42
43 tokio::fs::create_dir_all(dir.as_path())
44 .await
45 .map_err(|e| OpenAIError::FileSaveError(format!("{}, dir: {}", e, dir.display())))?;
46
47 tokio::fs::write(
48 file_path.as_path(),
49 response.bytes().await.map_err(|e| {
50 OpenAIError::FileSaveError(format!("{}, file path: {}", e, file_path.display()))
51 })?,
52 )
53 .await
54 .map_err(|e| OpenAIError::FileSaveError(e.to_string()))?;
55
56 Ok(file_path)
57}
58
59pub(crate) async fn save_b64<P: AsRef<Path>>(b64: &str, dir: P) -> Result<PathBuf, OpenAIError> {
60 let filename: String = rand::thread_rng()
61 .sample_iter(&Alphanumeric)
62 .take(10)
63 .map(char::from)
64 .collect();
65
66 let filename = format!("{filename}.png");
67
68 let path = PathBuf::from(dir.as_ref()).join(filename);
69
70 tokio::fs::write(
71 path.as_path(),
72 general_purpose::STANDARD
73 .decode(b64)
74 .map_err(|e| OpenAIError::FileSaveError(e.to_string()))?,
75 )
76 .await
77 .map_err(|e| OpenAIError::FileSaveError(format!("{}, path: {}", e, path.display())))?;
78
79 Ok(path)
80}