picox/src/storage_manager/mod.rs

use crate::model::image::{Image, ImageData, StorageLocation};
use crate::storage_manager::file_store::{FileStore, FileStoreConfig};
use axum::async_trait;
use j_db::database::Database;
use serde::{Deserialize, Serialize};
use std::fmt::{Display, Formatter};
use std::path::PathBuf;
use url::Url;

pub mod file_store;

#[derive(Debug)]
pub enum StoreError {
    InvalidFile,
    OutOfStorage,
    ImageTooBig,
    IOError(tokio::io::Error),
}

impl Display for StoreError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            StoreError::InvalidFile => write!(f, "Invalid file type"),
            StoreError::IOError(err) => write!(f, "IO Error: {}", err),
            StoreError::OutOfStorage => write!(f, "Underlying store full"),
            StoreError::ImageTooBig => write!(f, "Image too big for store"),
        }
    }
}

impl std::error::Error for StoreError {}

impl From<tokio::io::Error> for StoreError {
    fn from(value: tokio::io::Error) -> Self {
        Self::IOError(value)
    }
}

#[async_trait]
pub trait Store: Send {
    async fn store_img(
        &mut self,
        img_data: ImageData,
        file_name: &str,
    ) -> Result<(Url, StorageLocation), StoreError>;

    async fn create_img(
        &mut self,
        img_data: ImageData,
        file_name: &str,
        album: u64,
        created_by: u64,
    ) -> Result<Image, StoreError> {
        let (url, storage_location) = self.store_img(img_data, file_name).await?;

        Ok(Image::new(
            file_name,
            Vec::new(),
            url,
            created_by,
            storage_location,
            album,
        ))
    }

    async fn delete_img(&mut self, img: Image) -> Result<(), StoreError>;

    fn max_image_size(&self) -> usize;

    fn max_total_storage(&self) -> usize;

    async fn current_store_size(&self, db: &Database) -> Result<usize, StoreError>;
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum StorageTypes {
    FileStore,
    LinkStore,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageManagerConfig {
    pub default_storage_method: StorageTypes,
    pub allowed_types: Vec<String>,
    pub file_store_config: Option<FileStoreConfig>,
}

#[derive(Debug, Clone)]
pub struct StorageManager {
    config: StorageManagerConfig,

    file_store: Option<FileStore>,
}

impl StorageManager {
    pub fn new(storage_manager_config: StorageManagerConfig) -> Self {
        let file_store = storage_manager_config
            .file_store_config
            .clone()
            .map(FileStore::new);
        Self {
            config: storage_manager_config,
            file_store,
        }
    }

    fn get_default_storage_manager(&mut self) -> &mut dyn Store {
        self.get_storage_manager(self.config.default_storage_method)
    }

    fn get_storage_manager(&mut self, storage_type: StorageTypes) -> &mut dyn Store {
        match storage_type {
            StorageTypes::FileStore => self.file_store.as_mut().unwrap(),
            StorageTypes::LinkStore => {
                unimplemented!()
            }
        }
    }

    pub async fn store_img(
        &mut self,
        db: &Database,
        store: Option<StorageTypes>,
        img_data: ImageData,
        file_name: &str,
        created_by: u64,
        album: u64,
    ) -> Result<Image, StoreError> {
        let file_name_path = PathBuf::from(file_name);

        if let Some(ext) = file_name_path.extension() {
            let ext = ext.to_str().unwrap();
            if !self.config.allowed_types.contains(&ext.to_string()) {
                return Err(StoreError::InvalidFile);
            }
        } else {
            return Err(StoreError::InvalidFile);
        }

        let store_type = if let Some(store_type) = store {
            self.get_storage_manager(store_type)
        } else {
            self.get_default_storage_manager()
        };

        if img_data.size() > store_type.max_image_size() {
            return Err(StoreError::ImageTooBig);
        }

        let new_size = store_type.current_store_size(db).await? + img_data.size();

        if new_size > store_type.max_total_storage() {
            return Err(StoreError::OutOfStorage);
        }

        let img = store_type
            .create_img(img_data, file_name, created_by, album)
            .await?;

        Ok(db.insert::<Image>(img).unwrap())
    }
}