//! Tools for decrypting and loading collections from RPC responses.
use std::collections::HashMap;

use crate::{
    collection::{AnnotatedCollectionSchema, Collection},
    decoder::RecordSchema,
    indexer::{
        errors::LoadIndexerError,
        mapping::{FilterMapping, Mapping, MappingWithMeta, OreMapping},
        mapping_indexer::MappingIndexer,
        Indexer,
    },
};
use ciborium::{de::from_reader, ser::into_writer};
use cipherstash_grpc::api;
use envelopers::{DecryptionError, EncryptedRecord, EncryptionError, EnvelopeCipher, KeyProvider};
use rand_chacha::rand_core::{Error as RngError, RngCore};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use thiserror::Error;

#[derive(Error, Debug)]
pub enum DecryptIndexError {
    #[error("Failed to decrypt index settings: {0}")]
    DecryptSettings(#[from] DecryptionError),
    #[error("Failed to decode index settings from CBOR: {0}")]
    DecodeIndexSettings(#[from] ciborium::de::Error<std::io::Error>),
    #[error("{0}")]
    Other(String),
}

#[derive(Error, Debug)]
pub enum EncryptIndexError {
    #[error("Failed to encrypt index settings: {0}")]
    EncryptSettings(#[from] EncryptionError),
    #[error("Failed to encode index settings to CBOR: {0}")]
    EncodeIndexSettings(#[from] ciborium::ser::Error<std::io::Error>),
    #[error("{0}")]
    Other(String),
}

#[derive(Error, Debug)]
pub enum CreateIndexError {
    #[error("Failed to generate '{field}' key: {error}")]
    KeyGeneration {
        field: &'static str,
        error: RngError,
    },
}

#[derive(Error, Debug)]
pub enum DecryptCollectionMetadataError {
    #[error("Failed to decode metadata from CBOR: {0}")]
    Decode(#[from] ciborium::de::Error<std::io::Error>),
    #[error("Failed to decrypt metdata: {0}")]
    Decrypt(#[from] DecryptionError),
    #[error("{0}")]
    Other(String),
}

#[derive(Error, Debug)]
pub enum EncryptCollectionMetadataError {
    #[error("Failed to encode metadata to CBOR: {0}")]
    Encode(#[from] ciborium::ser::Error<std::io::Error>),
    #[error("Failed to decrypt metdata: {0}")]
    Encrypt(#[from] EncryptionError),
    #[error("{0}")]
    Other(String),
}

#[derive(Error, Debug)]
pub enum DecryptCollectionError {
    #[error("Failed to decrypt collection metadata: {0}")]
    DecryptMetadata(#[from] DecryptionError),
    #[error("Failed to decrypt index: {0}")]
    DecryptIndex(#[from] DecryptIndexError),
    #[error("Failed to decrypt metadata: {0}")]
    DecryptCollectionMetadata(#[from] DecryptCollectionMetadataError),
    #[error("{0}")]
    Other(String),
}

#[derive(Deserialize, Serialize, Debug)]
pub struct OreDecryptedIndexMeta {
    #[serde(rename = "$indexName")]
    index_name: String,
    #[serde(alias = "$prfKey", with = "hex")]
    prf_key: [u8; 16],
    #[serde(alias = "$prpKey", with = "hex")]
    prp_key: [u8; 16],
}

#[derive(Deserialize, Serialize, Debug)]
pub struct FilterDecryptedIndexMeta {
    #[serde(rename = "$indexName")]
    index_name: String,
    #[serde(alias = "$filterKey", with = "hex")]
    filter_key: [u8; 32],
}

#[derive(Deserialize, Serialize, Debug)]
#[serde(untagged)]
pub enum DecryptedIndexSettings {
    Ore {
        mapping: OreMapping,
        meta: OreDecryptedIndexMeta,
    },
    Filter {
        mapping: FilterMapping,
        meta: FilterDecryptedIndexMeta,
    },
}

impl DecryptedIndexSettings {
    pub fn index_name(&self) -> &str {
        match self {
            Self::Ore { meta, .. } => &meta.index_name,
            Self::Filter { meta, .. } => &meta.index_name,
        }
    }

    #[cfg(test)]
    fn to_mapping(self) -> Mapping {
        match self {
            Self::Ore { mapping, .. } => mapping.into(),
            Self::Filter { mapping, .. } => mapping.into(),
        }
    }

    pub fn into_mapping_with_meta(self, index_id: Uuid) -> MappingWithMeta {
        match self {
            Self::Ore {
                mapping,
                meta:
                    OreDecryptedIndexMeta {
                        prf_key, prp_key, ..
                    },
            } => MappingWithMeta::Ore {
                mapping,
                prf_key,
                prp_key,
                index_id,
            },
            Self::Filter {
                mapping,
                meta: FilterDecryptedIndexMeta { filter_key, .. },
            } => MappingWithMeta::Filter {
                mapping,
                filter_key,
                index_id,
            },
        }
    }
}

#[derive(Debug)]
#[allow(dead_code)]
pub struct DecryptedIndex {
    pub id: Uuid,
    pub settings: DecryptedIndexSettings,
    index_type: api::indexes::index::IndexType,
    unique: bool,
    first_schema_version: u32,
    last_schema_version: u32,
}

impl DecryptedIndex {
    /// Create a new index from a [`Mapping`] and index name.
    ///
    /// This method will generate encryption keys using the provided random number generator.
    pub(crate) fn new_from_mapping(
        rng: &mut impl RngCore,
        index_name: String,
        mapping: Mapping,
    ) -> Result<Self, CreateIndexError> {
        let unique = mapping.unique();
        let index_type = mapping.api_index_type();

        let settings = match mapping {
            Mapping::Ore(mapping) => {
                let mut prf_key = [0; 16];
                let mut prp_key = [0; 16];

                rng.try_fill_bytes(&mut prf_key).map_err(|error| {
                    CreateIndexError::KeyGeneration {
                        field: "prf_key",
                        error,
                    }
                })?;

                rng.try_fill_bytes(&mut prp_key).map_err(|error| {
                    CreateIndexError::KeyGeneration {
                        field: "prp_key",
                        error,
                    }
                })?;

                DecryptedIndexSettings::Ore {
                    mapping,
                    meta: OreDecryptedIndexMeta {
                        index_name,
                        prf_key,
                        prp_key,
                    },
                }
            }
            Mapping::Filter(mapping) => {
                let mut filter_key = [0; 32];

                rng.try_fill_bytes(&mut filter_key).map_err(|error| {
                    CreateIndexError::KeyGeneration {
                        field: "filter_key",
                        error,
                    }
                })?;

                DecryptedIndexSettings::Filter {
                    mapping,
                    meta: FilterDecryptedIndexMeta {
                        index_name,
                        filter_key,
                    },
                }
            }
        };

        Ok(Self {
            settings,
            id: Uuid::new_v4(),
            last_schema_version: 0,
            first_schema_version: 0,
            unique,
            index_type,
        })
    }

    /// Decrypt a collection index from an RPC Index
    async fn decrypt(
        cipher: &EnvelopeCipher<impl KeyProvider>,
        api::indexes::Index {
            id,
            settings,
            r#type,
            last_schema_version,
            first_schema_version,
            unique,
        }: api::indexes::Index,
    ) -> Result<Self, DecryptIndexError> {
        let id = id
            .try_into()
            .map(Uuid::from_bytes)
            .map_err(|_| DecryptIndexError::Other("Returned index id was malformed".into()))?;

        let decrypted_settings = cipher
            .decrypt(&EncryptedRecord::from_vec(settings).map_err(|_| {
                DecryptIndexError::Other(
                    "Failed to load encrypted index settings from bytes".into(),
                )
            })?)
            .await?;

        let settings = from_reader(&decrypted_settings[..])?;

        Ok(Self {
            id,
            settings,
            index_type: api::indexes::index::IndexType::from_i32(r#type)
                .unwrap_or(api::indexes::index::IndexType::Ore),
            last_schema_version,
            first_schema_version,
            unique,
        })
    }

    /// Encrypt the current index and convert it to a value that can be sent to the data service
    pub(crate) async fn encrypt(
        &self,
        cipher: &EnvelopeCipher<impl KeyProvider>,
    ) -> Result<api::indexes::Index, EncryptIndexError> {
        let mut settings_buffer = vec![];
        into_writer(&self.settings, &mut settings_buffer)?;
        let settings = cipher
            .encrypt(&settings_buffer)
            .await?
            .to_vec()
            .map_err(|e| EncryptIndexError::Other(e.to_string()))?;

        Ok(api::indexes::Index {
            id: self.id.as_bytes().to_vec(),
            settings,
            first_schema_version: self.first_schema_version,
            last_schema_version: self.last_schema_version,
            r#type: self.index_type.into(),
            unique: self.unique,
        })
    }

    /// Convert the decrypted index into a tuple in the form ( index_name, indexer ).
    ///
    /// These tuples could then be stored in a hash map like on the [`Indexer`] struct.
    fn into_indexer_entry(self) -> Result<(String, MappingIndexer), LoadIndexerError> {
        Ok((
            self.settings.index_name().into(),
            MappingIndexer::from_mapping(self.settings.into_mapping_with_meta(self.id))?,
        ))
    }
}

/// The decrypted metadata that is stored on the collection.
#[derive(Deserialize, Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct DecryptedCollectionMetadata {
    pub name: String,
    pub record_type: RecordSchema,
}

impl DecryptedCollectionMetadata {
    pub async fn decrypt(
        cipher: &EnvelopeCipher<impl KeyProvider>,
        encrypted_metadata: Vec<u8>,
    ) -> Result<Self, DecryptCollectionMetadataError> {
        let record = EncryptedRecord::from_vec(encrypted_metadata).map_err(|_| {
            DecryptCollectionMetadataError::Other(
                "Failed to load encrypted record from bytes".into(),
            )
        })?;

        let decrypted_bytes = cipher.decrypt(&record).await?;

        Ok(from_reader(&decrypted_bytes[..])?)
    }

    pub async fn encrypt(
        &self,
        cipher: &EnvelopeCipher<impl KeyProvider>,
    ) -> Result<Vec<u8>, EncryptCollectionMetadataError> {
        let mut bytes = vec![];

        into_writer(&self, &mut bytes)?;

        let encrypted_record = cipher.encrypt(&bytes).await?;

        encrypted_record.to_vec().map_err(|_| {
            EncryptCollectionMetadataError::Other(
                "Failed to encode encrypted record as bytes".into(),
            )
        })
    }
}

/// A collection that has been decrypted from an RPC [`InfoReply`]
#[derive(Debug)]
pub struct DecryptedCollection {
    pub id: Uuid,
    pub collection_ref: Vec<u8>,
    pub indexes: Vec<DecryptedIndex>,
    pub metadata: DecryptedCollectionMetadata,
}

impl DecryptedCollection {
    /// Load a collection from it's encrypted metadata and indexes
    pub(crate) async fn decrypt(
        cipher: &EnvelopeCipher<impl KeyProvider>,
        id: Uuid,
        collection_ref: Vec<u8>,
        encrypted_metadata: Vec<u8>,
        rpc_indexes: Vec<api::indexes::Index>,
    ) -> Result<Self, DecryptCollectionError> {
        let mut indexes = vec![];

        for index in rpc_indexes.into_iter() {
            indexes.push(DecryptedIndex::decrypt(cipher, index).await?);
        }

        let metadata = DecryptedCollectionMetadata::decrypt(cipher, encrypted_metadata).await?;

        Ok(Self {
            id,
            collection_ref,
            indexes,
            metadata,
        })
    }

    /// The plaintext name of the collection
    pub fn name(&self) -> &str {
        &self.metadata.name
    }
}

impl TryFrom<DecryptedCollection> for Collection {
    type Error = LoadIndexerError;

    fn try_from(col: DecryptedCollection) -> Result<Self, Self::Error> {
        let indexes = col
            .indexes
            .into_iter()
            .map(|x| x.into_indexer_entry())
            .collect::<Result<HashMap<_, _>, _>>()?;

        let indexer = Indexer::new(indexes, col.metadata.record_type);

        Ok(Self {
            id: col.id,
            name: col.metadata.name,
            collection_ref: col.collection_ref,
            indexer,
        })
    }
}

#[derive(Debug)]
pub enum CollectionDecryptResult {
    Decrypted(DecryptedCollection),
    DecryptionFailed {
        id: Uuid,
        r#ref: Vec<u8>,
        error: DecryptCollectionError,
    },
    Failed(DecryptCollectionError),
}

impl CollectionDecryptResult {
    /// Try and decrypt a [`CollectionDecryptResult`] from an RPC InfoReply.
    ///
    /// If the decryption fails a DecryptionFailed variant will be returned that contains the
    /// collection's ID. If the ID can't be decoded then a Failed variant will be returned.
    pub(crate) async fn decrypt(
        cipher: &EnvelopeCipher<impl KeyProvider>,
        api::collections::InfoReply {
            id,
            metadata,
            indexes,
            r#ref,
            ..
        }: api::collections::InfoReply,
    ) -> CollectionDecryptResult {
        let id_result = id.try_into().map(Uuid::from_bytes).map_err(|_| {
            DecryptCollectionError::Other("Returned collection id was malformed".into())
        });

        match id_result {
            Ok(id) => DecryptedCollection::decrypt(cipher, id, r#ref.clone(), metadata, indexes)
                .await
                .map_or_else(
                    |error| Self::DecryptionFailed { id, r#ref, error },
                    Self::Decrypted,
                ),
            Err(err) => Self::Failed(err),
        }
    }
}

impl TryFrom<CollectionDecryptResult> for Collection {
    type Error = DecryptCollectionError;

    fn try_from(value: CollectionDecryptResult) -> Result<Self, Self::Error> {
        match value {
            CollectionDecryptResult::Decrypted(collection) => collection
                .try_into()
                .map_err(|e: LoadIndexerError| DecryptCollectionError::Other(e.to_string())),
            CollectionDecryptResult::DecryptionFailed { error, .. } => Err(error),
            CollectionDecryptResult::Failed(error) => Err(error),
        }
    }
}

impl TryFrom<CollectionDecryptResult> for AnnotatedCollectionSchema {
    type Error = DecryptCollectionError;

    fn try_from(value: CollectionDecryptResult) -> Result<Self, Self::Error> {
        match value {
            CollectionDecryptResult::Decrypted(collection) => Ok(collection.into()),
            CollectionDecryptResult::DecryptionFailed { error, .. } => Err(error),
            CollectionDecryptResult::Failed(error) => Err(error),
        }
    }
}

#[cfg(test)]
mod test {
    use super::DecryptedCollection;
    use super::*;
    use crate::decoder::{DataType, RecordSchema, SchemaField};
    use crate::utils::{cbor_buffer, collection};
    use cipherstash_grpc::api::{collections::InfoReply, indexes::Index};
    use envelopers::{EnvelopeCipher, SimpleKeyProvider};
    use uuid::{uuid, Uuid};

    fn create_cipher() -> EnvelopeCipher<SimpleKeyProvider> {
        EnvelopeCipher::init(SimpleKeyProvider::init([1; 16]))
    }

    fn uuid_vec(val: &str) -> Vec<u8> {
        Uuid::parse_str(val).unwrap().as_bytes().to_vec()
    }

    async fn create_info_reply(cipher: &mut EnvelopeCipher<SimpleKeyProvider>) -> InfoReply {
        InfoReply {
            id: uuid_vec("33333c78-aaf4-488d-aef7-aeb5dc95b9e3"),
            metadata: cipher
                .encrypt(&cbor_buffer!({
                    "name" => "Great Collection",
                    "recordType" => {
                        "name" => "string",
                        "age" => "uint64"
                    }
                }))
                .await
                .unwrap()
                .to_vec()
                .unwrap(),
            indexes: vec![Index {
                id: uuid_vec("76a07734-a963-4b36-be5c-72cdfcdbf2eb"),
                settings: cipher
                    .encrypt(&cbor_buffer!({
                        "mapping" => {
                            "kind" => "exact",
                            "field" => "title"
                        },
                        "meta" => {
                            "$indexName" => "exactTitle",
                            "$prfKey" => hex::encode([0; 16]),
                            "$prpKey" => hex::encode([1; 16])
                        }
                    }))
                    .await
                    .unwrap()
                    .to_vec()
                    .unwrap(),
                ..Default::default()
            }],
            ..Default::default()
        }
    }

    #[tokio::test]
    async fn test_load_encrypted() {
        let mut cipher = create_cipher();

        let InfoReply {
            id,
            r#ref,
            metadata,
            indexes,
            ..
        } = create_info_reply(&mut cipher).await;

        let mut decrypted_collection = DecryptedCollection::decrypt(
            &mut cipher,
            Uuid::from_slice(&id).expect("Failed to load uuid"),
            r#ref,
            metadata,
            indexes,
        )
        .await
        .unwrap();

        assert_eq!(decrypted_collection.indexes.len(), 1);

        assert_eq!(decrypted_collection.name(), "Great Collection");

        assert_eq!(
            decrypted_collection.indexes.remove(0).settings.to_mapping(),
            OreMapping::Exact {
                field: "title".into(),
                unique: None
            }
            .into()
        );
    }

    #[tokio::test]
    async fn test_malformed_index_id() {
        let mut cipher = create_cipher();

        let InfoReply {
            id,
            r#ref,
            metadata,
            indexes,
            ..
        } = InfoReply {
            indexes: vec![Index {
                // invalid uuid
                id: vec![1; 10],
                ..Default::default()
            }],
            ..create_info_reply(&mut cipher).await
        };

        let error = DecryptedCollection::decrypt(
            &mut cipher,
            Uuid::from_slice(&id).expect("Failed to load uuid"),
            r#ref,
            metadata,
            indexes,
        )
        .await
        .unwrap_err();

        assert_eq!(
            format!("{}", error),
            "Failed to decrypt index: Returned index id was malformed"
        );
    }

    #[tokio::test]
    async fn test_info_reply_to_collection() {
        let mut cipher = create_cipher();

        let reply = create_info_reply(&mut cipher).await;

        let col: Collection = CollectionDecryptResult::decrypt(&mut cipher, reply)
            .await
            .try_into()
            .expect("Failed to convert result to collection");

        assert_eq!(col.name, "Great Collection");
    }

    #[tokio::test]
    async fn test_index_round_trip() {
        let cipher = create_cipher();

        let index = DecryptedIndex {
            id: uuid!("be4edac0-a692-4689-b3a2-584590ee2c2c"),
            settings: DecryptedIndexSettings::Ore {
                mapping: OreMapping::Exact {
                    field: "test".into(),
                    unique: None,
                },
                meta: OreDecryptedIndexMeta {
                    index_name: "great index".into(),
                    prf_key: [0; 16],
                    prp_key: [1; 16],
                },
            },
            unique: false,
            index_type: api::indexes::index::IndexType::Ore,
            first_schema_version: 0,
            last_schema_version: 0,
        };

        let encrypted = index.encrypt(&cipher).await.expect("Failed to encrypt");

        let decrypted = DecryptedIndex::decrypt(&cipher, encrypted)
            .await
            .expect("Failed to encrypt");

        assert_eq!(decrypted.id, index.id);
        assert_eq!(decrypted.settings.index_name(), "great index");
        assert_eq!(decrypted.unique, false);
        assert_eq!(
            decrypted.settings.to_mapping(),
            OreMapping::Exact {
                field: "test".into(),
                unique: None
            }
            .into()
        );
    }

    #[tokio::test]
    async fn test_collection_metadata_round_trip() {
        let cipher = create_cipher();

        let metadata = DecryptedCollectionMetadata {
            name: "Great Metadata".into(),
            record_type: RecordSchema {
                map: collection! {
                    "uint64" => DataType::Uint64,
                    "bool" => DataType::Boolean
                },
            },
        };

        let encrypted = metadata.encrypt(&cipher).await.expect("Failed to encrypt");

        let decrypted = DecryptedCollectionMetadata::decrypt(&cipher, encrypted)
            .await
            .expect("Failed to encrypt");

        assert_eq!(decrypted.name, "Great Metadata");
        assert_eq!(
            decrypted.record_type.map.get("uint64"),
            Some(&SchemaField::DataType(DataType::Uint64))
        );
        assert_eq!(
            decrypted.record_type.map.get("bool"),
            Some(&SchemaField::DataType(DataType::Boolean))
        );
    }
}