/*
 * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
 *
 * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
 */

use crate::store::{CONFIG, TempDir, cleanup::store_destroy};
use ahash::AHashMap;
use common::{Core, Inner, Server, config::storage::Storage};
use email::message::metadata::MessageMetadata;
use http::management::stores::destroy_account_blobs;
use std::sync::Arc;
use store::{
    BlobStore, Serialize, SerializeInfallible, Stores,
    write::{Archiver, BatchBuilder, BlobLink, BlobOp, ValueClass, blob::BlobQuota, now},
};
use types::{blob::BlobClass, blob_hash::BlobHash, collection::Collection, field::EmailField};
use utils::config::Config;

#[tokio::test]
pub async fn blob_tests() {
    let temp_dir = TempDir::new("blob_tests", true);
    let mut config =
        Config::new(CONFIG.replace("{TMP}", temp_dir.path.as_path().to_str().unwrap())).unwrap();
    let stores = Stores::parse_all(&mut config, false).await;

    for (store_id, blob_store) in &stores.blob_stores {
        println!("Testing blob store {}...", store_id);
        test_store(blob_store.clone()).await;
    }

    for (store_id, store) in stores.stores {
        println!("Testing blob management on store {}...", store_id);

        // Init store
        store_destroy(&store).await;

        // Test internal blob store
        let blob_store: BlobStore = store.clone().into();
        let server = Server {
            inner: Arc::new(Inner::default()),
            core: Arc::new(Core {
                storage: Storage {
                    data: store.clone(),
                    blob: blob_store.clone(),
                    ..Default::default()
                },
                ..Default::default()
            }),
        };

        // Blob hash exists
        let hash = BlobHash::generate(b"abc".as_slice());
        assert!(!store.blob_exists(&hash).await.unwrap());

        // Reserve blob
        let until = now() + 1;
        store
            .write(
                BatchBuilder::new()
                    .with_account_id(0)
                    .set(
                        BlobOp::Link {
                            to: BlobLink::Temporary { until },
                            hash: hash.clone(),
                        },
                        1024u32.serialize(),
                    )
                    .build_all(),
            )
            .await
            .unwrap();

        // Uncommitted blob, should not exist
        assert!(!store.blob_exists(&hash).await.unwrap());

        // Write blob to store
        blob_store.put_blob(hash.as_ref(), b"abc").await.unwrap();

        // Commit blob
        store
            .write(
                BatchBuilder::new()
                    .set(BlobOp::Commit { hash: hash.clone() }, Vec::new())
                    .build_all(),
            )
            .await
            .unwrap();

        // Blob hash should now exist
        assert!(store.blob_exists(&hash).await.unwrap());
        assert!(
            blob_store
                .get_blob(hash.as_ref(), 0..usize::MAX)
                .await
                .unwrap()
                .is_some()
        );

        // AccountId 0 should be able to read blob
        assert!(
            store
                .blob_has_access(
                    &hash,
                    BlobClass::Reserved {
                        account_id: 0,
                        expires: until
                    }
                )
                .await
                .unwrap()
        );

        // AccountId 1 should not be able to read blob
        assert!(
            !store
                .blob_has_access(
                    &hash,
                    BlobClass::Reserved {
                        account_id: 1,
                        expires: until
                    }
                )
                .await
                .unwrap()
        );

        // Blob already expired, quota should be 0
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        assert_eq!(
            store.blob_quota(0).await.unwrap(),
            BlobQuota { bytes: 0, count: 0 }
        );

        // Purge expired blobs
        store.purge_blobs(blob_store.clone()).await.unwrap();

        // Blob hash should no longer exist
        assert!(!store.blob_exists(&hash).await.unwrap());

        // AccountId 0 should not be able to read blob
        assert!(
            !store
                .blob_has_access(
                    &hash,
                    BlobClass::Reserved {
                        account_id: 0,
                        expires: until
                    }
                )
                .await
                .unwrap()
        );

        // Blob should no longer be in store
        assert!(
            blob_store
                .get_blob(hash.as_ref(), 0..usize::MAX)
                .await
                .unwrap()
                .is_none()
        );

        // Upload one linked blob to accountId 1, two linked blobs to accountId 0, and three unlinked (reserved) blobs to accountId 2
        let expiry_times = AHashMap::from_iter([
            (b"abc", now() - 10),
            (b"efg", now() + 10),
            (b"hij", now() + 10),
        ]);
        for (document_id, (blob, blob_value)) in [
            (b"123", vec![]),
            (b"456", vec![]),
            (b"789", vec![]),
            (b"abc", 5000u32.serialize()),
            (b"efg", 1000u32.serialize()),
            (b"hij", 2000u32.serialize()),
        ]
        .into_iter()
        .enumerate()
        {
            let hash = BlobHash::generate(blob.as_slice());
            let mut batch = BatchBuilder::new();
            batch
                .with_account_id(if document_id > 0 { 0 } else { 1 })
                .with_collection(Collection::Email)
                .with_document(document_id as u32);
            if let Some(until) = expiry_times.get(blob) {
                if !blob_value.is_empty() {
                    batch.set(
                        BlobOp::Quota {
                            hash: hash.clone(),
                            until: *until,
                        },
                        blob_value,
                    );
                }
                batch.set(
                    BlobOp::Link {
                        hash: hash.clone(),
                        to: BlobLink::Temporary { until: *until },
                    },
                    vec![],
                );
            } else {
                batch
                    .set(
                        BlobOp::Link {
                            hash: hash.clone(),
                            to: BlobLink::Document,
                        },
                        vec![],
                    )
                    .set(
                        ValueClass::Property(EmailField::Metadata.into()),
                        Archiver::new(MessageMetadata {
                            contents: Default::default(),
                            rcvd_attach: Default::default(),
                            blob_hash: hash.clone(),
                            blob_body_offset: Default::default(),
                            preview: Default::default(),
                            raw_headers: Default::default(),
                        })
                        .serialize()
                        .unwrap(),
                    );
            };
            batch.set(BlobOp::Commit { hash: hash.clone() }, vec![]);

            store.write(batch.build_all()).await.unwrap();
            blob_store
                .put_blob(hash.as_ref(), blob.as_slice())
                .await
                .unwrap();
        }

        // One of the reserved blobs expired and should not count towards quota
        assert_eq!(
            store.blob_quota(0).await.unwrap(),
            BlobQuota {
                bytes: 3000,
                count: 2
            }
        );
        assert_eq!(
            store.blob_quota(1).await.unwrap(),
            BlobQuota { bytes: 0, count: 0 }
        );

        // Purge expired blobs and make sure nothing else is deleted
        store.purge_blobs(blob_store.clone()).await.unwrap();
        for (pos, (blob, blob_class)) in [
            (
                b"abc",
                BlobClass::Reserved {
                    account_id: 0,
                    expires: expiry_times[&b"abc"],
                },
            ),
            (
                b"123",
                BlobClass::Linked {
                    account_id: 1,
                    collection: 0,
                    document_id: 0,
                },
            ),
            (
                b"456",
                BlobClass::Linked {
                    account_id: 0,
                    collection: 0,
                    document_id: 1,
                },
            ),
            (
                b"789",
                BlobClass::Linked {
                    account_id: 0,
                    collection: 0,
                    document_id: 2,
                },
            ),
            (
                b"efg",
                BlobClass::Reserved {
                    account_id: 0,
                    expires: expiry_times[&b"efg"],
                },
            ),
            (
                b"hij",
                BlobClass::Reserved {
                    account_id: 0,
                    expires: expiry_times[&b"hij"],
                },
            ),
        ]
        .into_iter()
        .enumerate()
        {
            let ct = pos == 0;
            let hash = BlobHash::generate(blob.as_slice());
            assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
            assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
            assert!(
                blob_store
                    .get_blob(hash.as_ref(), 0..usize::MAX)
                    .await
                    .unwrap()
                    .is_some()
                    ^ ct
            );
        }

        // AccountId 0 should not have access to accountId 1's blobs
        assert!(
            !store
                .blob_has_access(
                    BlobHash::generate(b"123".as_slice()),
                    BlobClass::Linked {
                        account_id: 0,
                        collection: 0,
                        document_id: 0,
                    }
                )
                .await
                .unwrap()
        );

        // Unlink blob
        store
            .write(
                BatchBuilder::new()
                    .with_account_id(0)
                    .with_collection(Collection::Email)
                    .with_document(2)
                    .clear(BlobOp::Link {
                        hash: BlobHash::generate(b"789".as_slice()),
                        to: BlobLink::Document,
                    })
                    .build_all(),
            )
            .await
            .unwrap();

        // Purge and make sure blob is deleted
        store.purge_blobs(blob_store.clone()).await.unwrap();
        for (pos, (blob, blob_class)) in [
            (
                b"789",
                BlobClass::Linked {
                    account_id: 0,
                    collection: 0,
                    document_id: 2,
                },
            ),
            (
                b"123",
                BlobClass::Linked {
                    account_id: 1,
                    collection: 0,
                    document_id: 0,
                },
            ),
            (
                b"456",
                BlobClass::Linked {
                    account_id: 0,
                    collection: 0,
                    document_id: 1,
                },
            ),
            (
                b"efg",
                BlobClass::Reserved {
                    account_id: 0,
                    expires: expiry_times[&b"efg"],
                },
            ),
            (
                b"hij",
                BlobClass::Reserved {
                    account_id: 0,
                    expires: expiry_times[&b"hij"],
                },
            ),
        ]
        .into_iter()
        .enumerate()
        {
            let ct = pos == 0;
            let hash = BlobHash::generate(blob.as_slice());
            assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
            assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
            assert!(
                blob_store
                    .get_blob(hash.as_ref(), 0..usize::MAX)
                    .await
                    .unwrap()
                    .is_some()
                    ^ ct
            );
        }

        // Unlink all blobs from accountId 1 and purge
        destroy_account_blobs(&server, 1).await.unwrap();
        store.purge_blobs(blob_store.clone()).await.unwrap();

        // Make sure only accountId 0's blobs are left
        for (pos, (blob, blob_class)) in [
            (
                b"123",
                BlobClass::Linked {
                    account_id: 1,
                    collection: 0,
                    document_id: 0,
                },
            ),
            (
                b"456",
                BlobClass::Linked {
                    account_id: 0,
                    collection: 0,
                    document_id: 1,
                },
            ),
            (
                b"efg",
                BlobClass::Reserved {
                    account_id: 0,
                    expires: expiry_times[&b"efg"],
                },
            ),
            (
                b"hij",
                BlobClass::Reserved {
                    account_id: 0,
                    expires: expiry_times[&b"hij"],
                },
            ),
        ]
        .into_iter()
        .enumerate()
        {
            let ct = pos == 0;
            let hash = BlobHash::generate(blob.as_slice());
            assert!(store.blob_has_access(&hash, blob_class).await.unwrap() ^ ct);
            assert!(store.blob_exists(&hash).await.unwrap() ^ ct);
            assert!(
                blob_store
                    .get_blob(hash.as_ref(), 0..usize::MAX)
                    .await
                    .unwrap()
                    .is_some()
                    ^ ct
            );
        }
    }
    temp_dir.delete();
}

async fn test_store(store: BlobStore) {
    // Test small blob
    const DATA: &[u8] = b"Lorem ipsum dolor sit amet, consectetur adipiscing elit. Fusce erat nisl, dignissim a porttitor id, varius nec arcu. Sed mauris.";
    let hash = BlobHash::generate(DATA);

    store.put_blob(hash.as_slice(), DATA).await.unwrap();
    assert_eq!(
        String::from_utf8(
            store
                .get_blob(hash.as_slice(), 0..usize::MAX)
                .await
                .unwrap()
                .unwrap()
        )
        .unwrap(),
        std::str::from_utf8(DATA).unwrap()
    );
    assert_eq!(
        String::from_utf8(
            store
                .get_blob(hash.as_slice(), 11..57)
                .await
                .unwrap()
                .unwrap()
        )
        .unwrap(),
        std::str::from_utf8(&DATA[11..57]).unwrap()
    );
    assert!(store.delete_blob(hash.as_slice()).await.unwrap());
    assert!(
        store
            .get_blob(hash.as_slice(), 0..usize::MAX)
            .await
            .unwrap()
            .is_none()
    );

    // Test large blob
    let mut data = Vec::with_capacity(50 * 1024 * 1024);
    while data.len() < 50 * 1024 * 1024 {
        data.extend_from_slice(DATA);
        let marker = format!(" [{}] ", data.len());
        data.extend_from_slice(marker.as_bytes());
    }
    let hash = BlobHash::generate(&data);
    store.put_blob(hash.as_slice(), &data).await.unwrap();
    assert_eq!(
        String::from_utf8(
            store
                .get_blob(hash.as_slice(), 0..usize::MAX)
                .await
                .unwrap()
                .unwrap()
        )
        .unwrap(),
        std::str::from_utf8(&data).unwrap()
    );

    assert_eq!(
        String::from_utf8(
            store
                .get_blob(hash.as_slice(), 3000111..4000999)
                .await
                .unwrap()
                .unwrap()
        )
        .unwrap(),
        std::str::from_utf8(&data[3000111..4000999]).unwrap()
    );
    assert!(store.delete_blob(hash.as_slice()).await.unwrap());
    assert!(
        store
            .get_blob(hash.as_slice(), 0..usize::MAX)
            .await
            .unwrap()
            .is_none()
    );
}
