use std::pin::Pin;

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

use crate::config::MixedStorageConfig;

use super::{disk::DiskStorage, memory::MemoryStorage, Storage};

pub struct MixedStorage {
    disk: DiskStorage,
    memory: MemoryStorage,
}

impl MixedStorage {
    pub fn new(config: MixedStorageConfig) -> Self {
        Self {
            disk: DiskStorage::new(config.clone().into()),
            memory: MemoryStorage::new(config.into()),
        }
    }
}

#[async_trait]
impl Storage for MixedStorage {
    async fn init(&mut self) -> anyhow::Result<()> {
        self.disk.init().await?;
        self.memory.init().await?;

        let files = self.disk.retrieve_all().await?;
        for (key, stream) in files.into_iter() {
            self.memory.save(key, stream).await?;
        }

        Ok(())
    }

    async fn tick(&mut self) -> anyhow::Result<()> {
        Ok(())
    }

    async fn eligible(&self, _src: String) -> bool {
        true
    }

    async fn delete(&mut self, _key: String) -> anyhow::Result<()> {
        Ok(())
    }

    async fn retrieve(
        &self,
        key: String,
    ) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>> {
        self.memory.retrieve(key).await
    }

    async fn save(
        &mut self,
        key: String,
        stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
    ) -> anyhow::Result<()> {
        self.memory.save(key.clone(), stream).await?;

        let final_stream = self.memory.retrieve(key.clone()).await.unwrap();
        self.disk.save(key, final_stream).await?;

        Ok(())
    }
}