scel/crates/scel_core/src/services/mod.rs

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<String>,
    pub link: String,
    pub progress: Option<u32>,
    pub file_name: Option<String>,
}

pub struct InMemoryDownloadService {
    downloads: Mutex<
        HashMap<
            String,
            (
                Arc<Mutex<Download>>,
                Arc<Mutex<tokio::sync::watch::Sender<Download>>>,
                tokio::sync::watch::Receiver<Download>,
            ),
        >,
    >,
}

impl InMemoryDownloadService {
    pub fn new() -> Self {
        Self {
            downloads: Mutex::new(HashMap::new()),
        }
    }

    pub async fn add_download(self: Arc<Self>, download: Download) -> anyhow::Result<Download> {
        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<Self>, 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<Option<Download>> {
        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<Download> {
        let downloads = self.downloads.lock().await;
        let download = downloads.get(&id).unwrap();
        download.2.clone()
    }
}