pub mod config;
mod local;
mod s3;

use std::pin::Pin;
use std::sync::Arc;

use crate::local::LocalStorage;
use crate::s3::S3Storage;
use async_trait::async_trait;
use bytes::Bytes;
use thiserror::Error;
use tokio::io::AsyncBufRead;
use tokio::sync::mpsc;

#[derive(Error, Debug)]
pub enum BlobStorageError {
    #[error("S3 Download error: {0}")]
    S3DownloadSDKError(
        #[from] aws_sdk_s3::error::SdkError<aws_sdk_s3::operation::get_object::GetObjectError>,
    ),
    #[error("S3 Upload error: {0}")]
    S3UploadFileSDKError(
        #[from] aws_sdk_s3::error::SdkError<aws_sdk_s3::operation::put_object::PutObjectError>,
    ),
    #[error("S3 Upload part error: {0}")]
    S3UploadPartSDKError(
        #[from] aws_sdk_s3::error::SdkError<aws_sdk_s3::operation::upload_part::UploadPartError>,
    ),
    #[error("S3 Delete error: {0}")]
    S3DeleteSDKError(
        #[from]
        aws_sdk_s3::error::SdkError<aws_sdk_s3::operation::delete_object::DeleteObjectError>,
    ),
    #[error("S3 Complete multipart upload error: {0}")]
    S3CompleteMultipartUploadSDKError(
        #[from]
        aws_sdk_s3::error::SdkError<
            aws_sdk_s3::operation::complete_multipart_upload::CompleteMultipartUploadError,
        >,
    ),
    #[error("S3 Initiate multipart upload error: {0}")]
    S3InitiateMultipartUploadSDKError(
        #[from]
        aws_sdk_s3::error::SdkError<
            aws_sdk_s3::operation::create_multipart_upload::CreateMultipartUploadError,
        >,
    ),
    #[error("S3 No upload_id returned")]
    S3NoUploadIdReturned,
    #[error("S3 No ETag returned")]
    S3NoETagReturned,
    #[error("Local Download error: {0}")]
    DownloadLocalError(#[from] std::io::Error),
    #[error("No storage configuration found")]
    NoStorageConfiguration,
    #[error("Failed to parse configuration: {0}")]
    ConfigurationError(#[from] ::config::ConfigError),
    #[error("Failed to send data to channel")]
    ChannelSendError(#[from] tokio::sync::mpsc::error::SendError<FileReader>),
}

pub const PART_SIZE: usize = 5 * 1024 * 1024;

/// A stream of bytes, typically from a file or HTTP body
pub type ByteStream = Pin<Box<dyn AsyncBufRead + Send + Sync>>;

pub struct FileReader {
    pub key: String,
    pub data: ByteStream,
}

pub struct File {
    pub key: String,
    pub data: Bytes,
}

#[async_trait]
pub trait BlobStorage: Send + Sync + 'static {
    async fn download_files(
        &self,
        keys: &[&str],
    ) -> Result<mpsc::Receiver<FileReader>, BlobStorageError>;
    async fn download_file(&self, key: &str) -> Result<FileReader, BlobStorageError>;
    async fn upload_file(&self, file: File) -> Result<Option<String>, BlobStorageError>;
    async fn stream_upload(
        &self,
        key: &str,
        reader: ByteStream,
    ) -> Result<Option<String>, BlobStorageError>;
    async fn upload_files(
        &self,
        files: mpsc::Receiver<File>,
    ) -> Result<Vec<String>, BlobStorageError>;
    async fn delete_file(&self, key: &str) -> Result<(), BlobStorageError>;
}

pub async fn create_storage() -> Result<Arc<dyn BlobStorage>, BlobStorageError> {
    let config = config::Config::from_env()?;
    if let Some(s3_config) = &config.s3_config {
        let storage = S3Storage::new(s3_config);
        Ok(Arc::new(storage))
    } else if let Some(local_config) = &config.local_config {
        let storage = LocalStorage::new(local_config);
        Ok(Arc::new(storage))
    } else {
        Err(BlobStorageError::NoStorageConfiguration)
    }
}


#[cfg(test)]
mod tests {
    use crate::create_storage;

    use super::*;
    use std::sync::Arc;
    use bytes::Bytes;
    use tokio::io::{AsyncReadExt, BufStream};

    async fn setup_storage() -> Vec<Arc<dyn BlobStorage>> {
        let s3_storage = {
            unsafe {std::env::set_var("BLOB_STORAGE_TYPE", "s3");}
            dotenvy::dotenv().ok();
            create_storage().await.expect("create s3 storage failed")
        };
        let local_storage = {
            unsafe {std::env::set_var("BLOB_STORAGE_TYPE", "local");}
            dotenvy::dotenv().ok();
            create_storage().await.expect("create local storage failed")
        };
        vec![s3_storage, local_storage]
    }

    // Create a random byte vector
    fn make_png_vec() -> Vec<u8> {
        vec![
            0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG header
            0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00,
            0x00, 0x01, 0x08, 0x06, 0x00, 0x00, 0x00, 0x1F, 0x15, 0xC4, 0x89, 0x00, 0x00, 0x00,
            0x0A, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9C, 0x63, 0x00, 0x01, 0x00, 0x00, 0x05, 0x00,
            0x01, 0x0D, 0x0A, 0x2D, 0xB4, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE,
            0x42, 0x60, 0x82,
        ]
    }

    fn make_byte_stream(size: usize) -> ByteStream {
        let mut data = Vec::new();
        for _ in 0..size {
            data.push(rand::random::<u8>());
        }
        let cursor = std::io::Cursor::new(data);
        let buf_reader = BufStream::new(cursor);
        Box::pin(buf_reader)
    }

    #[tokio::test]
    async fn test_upload_file_and_get_url() {
        for storage in setup_storage().await {
            let test_data = make_png_vec();

            let key = "test/test_upload_file_and_get_url.png";

            // Upload file
            let url = storage.upload_file(File {
                key: key.to_string(),
                data: Bytes::from(test_data.clone()),
            }).await.expect("upload_file failed");
            assert_eq!(url, None);

            // Get content, ensure it matches
            let mut stream = storage.download_file(key).await.expect("download_file failed");
            let mut buf = Vec::new();
            stream.data.read_to_end(&mut buf).await.expect("read_to_end error");
            assert_eq!(buf, test_data);

            // Clean up
            storage.delete_file(key).await.expect("delete_file failed");
        }
    }

    #[tokio::test]
    async fn test_multipart_upload_large_file() {
        for storage in setup_storage().await {
            let key = "test/test_multipart_upload_large_file.png";

            // Upload file
            let _url = storage.stream_upload(key, make_byte_stream(2*PART_SIZE)).await.expect("stream_upload failed with larger chunk");

            // Get content, ensure it matches
            let mut stream = storage.download_file(key).await.expect("download_file failed");
            let mut buf = Vec::new();
            stream.data.read_to_end(&mut buf).await.expect("read_to_end error");
            assert_eq!(buf.len(), PART_SIZE*2);

            // Clean up
            storage.delete_file(key).await.expect("delete_file failed");

            let _url = storage.stream_upload(key, make_byte_stream(PART_SIZE)).await.expect("stream_upload failed with PART_SIZE chunk");
            let mut stream = storage.download_file(key).await.expect("download_file failed");
            let mut buf = Vec::new();
            stream.data.read_to_end(&mut buf).await.expect("read_to_end error");
            assert_eq!(buf.len(), PART_SIZE);

            // Clean up
            storage.delete_file(key).await.expect("delete_file failed");

            let _url = storage.stream_upload(key, make_byte_stream(PART_SIZE/2)).await.expect("stream_upload failed with smaller chunk");
            let mut stream = storage.download_file(key).await.expect("download_file failed");
            let mut buf = Vec::new();
            stream.data.read_to_end(&mut buf).await.expect("read_to_end error");
            assert_eq!(buf.len(), PART_SIZE/2);

            // Clean up
            storage.delete_file(key).await.expect("delete_file failed");
        }
    }
}
