Kordophone/kordophoned/src/daemon/attachment_store.rs

use std::{
    io::{BufWriter, Write},
    path::PathBuf,
};

use anyhow::Result;
use futures_util::StreamExt;
use kordophone::APIInterface;
use thiserror::Error;
use tokio::pin;

mod target {
    pub static ATTACHMENTS: &str = "attachments";
}

#[derive(Debug, Clone)]
pub struct Attachment {
    pub guid: String,
    pub path: PathBuf,
    pub downloaded: bool,
}

#[derive(Debug, Error)]
enum AttachmentStoreError {
    #[error("attachment has already been downloaded")]
    AttachmentAlreadyDownloaded,

    #[error("Client error: {0}")]
    APIClientError(String),
}

pub struct AttachmentStore {
    store_path: PathBuf,
}

impl AttachmentStore {
    pub fn new(data_dir: &PathBuf) -> AttachmentStore {
        let store_path = data_dir.join("attachments");
        log::info!(target: target::ATTACHMENTS, "Attachment store path: {}", store_path.display());

        // Create the attachment store if it doesn't exist
        std::fs::create_dir_all(&store_path)
            .expect("Wasn't able to create the attachment store path");

        AttachmentStore {
            store_path: store_path,
        }
    }

    pub fn get_attachment(&self, guid: &String) -> Attachment {
        let path = self.store_path.join(guid);
        let path_exists = std::fs::exists(&path).expect(
            format!(
                "Wasn't able to check for the existence of an attachment file path at {}",
                &path.display()
            )
            .as_str(),
        );

        Attachment {
            guid: guid.to_owned(),
            path: path,
            downloaded: path_exists,
        }
    }

    pub async fn download_attachment<C, F, Fut>(
        &mut self,
        attachment: &Attachment,
        mut client_factory: F,
    ) -> Result<()>
    where
        C: APIInterface,
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = Result<C>>,
    {
        if attachment.downloaded {
            log::info!(target: target::ATTACHMENTS, "Attachment already downloaded: {}", attachment.guid);
            return Err(AttachmentStoreError::AttachmentAlreadyDownloaded.into());
        }

        log::info!(target: target::ATTACHMENTS, "Starting download for attachment: {}", attachment.guid);

        // Create temporary file first, we'll atomically swap later.
        assert!(!std::fs::exists(&attachment.path).unwrap());
        let file = std::fs::File::create(&attachment.path)?;
        let mut writer = BufWriter::new(&file);

        log::trace!(target: target::ATTACHMENTS, "Created attachment file at {}", &attachment.path.display());

        let mut client = client_factory().await?;
        let stream = client
            .fetch_attachment_data(&attachment.guid)
            .await
            .map_err(|e| AttachmentStoreError::APIClientError(format!("{:?}", e)))?;

        // Since we're async, we need to pin this.
        pin!(stream);

        log::trace!(target: target::ATTACHMENTS, "Writing attachment data to disk");
        while let Some(Ok(data)) = stream.next().await {
            writer.write(data.as_ref())?;
        }

        log::info!(target: target::ATTACHMENTS, "Completed download for attachment: {}", attachment.guid);
        Ok(())
    }
    
    /// Check if an attachment should be downloaded
    pub fn should_download(&self, attachment_id: &str) -> bool {
        let attachment = self.get_attachment(&attachment_id.to_string());
        !attachment.downloaded
    }
}