use std::{collections::HashMap, path::PathBuf, sync::Arc}; use tokio::sync::{watch, Mutex}; use tracing::error; use uuid::Uuid; use crate::youtube::{Arg, YoutubeDL}; #[derive(Clone)] pub struct Download { pub id: Option, pub link: String, pub progress: Option, pub file_name: Option, } pub struct InMemoryDownloadService { downloads: Mutex< HashMap< String, ( Arc>, Arc>>, tokio::sync::watch::Receiver, ), >, >, } impl InMemoryDownloadService { pub fn new() -> Self { Self { downloads: Mutex::new(HashMap::new()), } } pub async fn add_download(self: Arc, download: Download) -> anyhow::Result { let mut downloads = self.downloads.lock().await; let (tx, rx) = watch::channel(download.clone()); let shared_tx = Arc::new(Mutex::new(tx)); let mut d = download.to_owned(); let id = Uuid::new_v4().to_string(); d.id = Some(id.clone()); downloads.insert(id.clone(), (Arc::new(Mutex::new(d.clone())), shared_tx, rx)); let args = vec![ Arg::new("--progress"), Arg::new("--newline"), Arg::new_with_args("--output", "%(title).90s.%(ext)s"), ]; let ytd = YoutubeDL::new( &PathBuf::from("./data/downloads"), args, download.link.as_str(), )?; tokio::spawn({ let download_service = self.clone(); async move { if let Err(e) = ytd .download( |percentage| { let ds = download_service.clone(); let id = id.clone(); async move { let mut download = ds.get_download(id).await.unwrap().unwrap(); download.progress = Some(percentage); let _ = ds.update_download(download).await; } }, |file_name| { let ds = download_service.clone(); let id = id.clone(); async move { let mut download = ds.get_download(id).await.unwrap().unwrap(); download.file_name = Some(file_name); let _ = ds.update_download(download).await; } }, ) .await { error!("Download failed: {}", e); } else { let download = download_service.get_download(id).await.unwrap().unwrap(); let _ = download_service.update_download(download).await; } } }); Ok(d) } pub async fn update_download(self: Arc, download: Download) -> anyhow::Result<()> { let mut downloads = self.downloads.lock().await; if let Some(d) = downloads.get_mut(&download.clone().id.unwrap()) { let mut d_mut = d.0.lock().await; *d_mut = download.clone(); let _ = d.1.lock().await.send(download); } Ok(()) } pub async fn get_download(&self, id: String) -> anyhow::Result> { let downloads = self.downloads.lock().await; if let Some(d) = downloads.get(&id) { let download = d.0.lock().await; Ok(Some(download.clone())) } else { Ok(None) } } pub async fn subscribe_download(&self, id: String) -> tokio::sync::watch::Receiver { let downloads = self.downloads.lock().await; let download = downloads.get(&id).unwrap(); download.2.clone() } }