pub mod disk;
pub mod memory;
pub mod mixed;

use std::pin::Pin;

use async_trait::async_trait;
use tokio_stream::Stream;
use tokio_util::bytes::Bytes;

use crate::{
    config::{StorageConfig, StorageStrategyConfig},
    crypt,
};

use self::{disk::DiskStorage, memory::MemoryStorage, mixed::MixedStorage};

#[async_trait]
pub trait Storage {
    async fn init(&mut self) -> anyhow::Result<()>;
    async fn eligible(&self, src: String) -> bool;
    async fn delete(&mut self, key: String) -> anyhow::Result<()>;
    async fn retrieve(
        &self,
        key: String,
    ) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>>;
    async fn save(
        &mut self,
        key: String,
        stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
    ) -> anyhow::Result<()>;
}

pub struct StoragePool {
    storages: Vec<Box<dyn Storage + Sync + Send>>,
}

impl StoragePool {
    pub fn from_config(items: Vec<StorageConfig>) -> Self {
        Self {
            storages: items.into_iter().map(create_storage).collect(),
        }
    }

    pub async fn init(&mut self) -> anyhow::Result<()> {
        for item in &mut self.storages {
            item.init().await?;
        }

        Ok(())
    }

    pub async fn retrieve(
        &self,
        src: String,
    ) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>> {
        let mut stream = None;
        let key = crypt::compute_key(src);

        for item in &self.storages {
            if let Some(storage_stream) = item.retrieve(key.clone()).await {
                stream = Some(storage_stream);
                break;
            }
        }

        stream
    }

    pub async fn save(
        &mut self,
        src: String,
        stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
    ) -> anyhow::Result<()> {
        let key = crypt::compute_key(src);

        for item in self.storages.iter_mut() {
            if item.eligible(key.clone()).await {
                item.save(key, stream).await?;
                break;
            }
        }

        Ok(())
    }
}

fn create_storage(item: StorageConfig) -> Box<dyn Storage + Sync + Send> {
    match item.strategy {
        StorageStrategyConfig::Memory(config) => Box::new(MemoryStorage::new(config)),
        StorageStrategyConfig::Disk(config) => Box::new(DiskStorage::new(config)),
        StorageStrategyConfig::Mixed(config) => Box::new(MixedStorage::new(config)),
    }
}