pub(crate) mod helpers;
pub mod value;

use cipherstash_grpc::api::documents::Document;
use envelopers::{DecryptionError, EncryptedRecord, EncryptionError, EnvelopeCipher, KeyProvider};
use uuid::Uuid;
pub use value::Value;

use minicbor::{Decoder, Encoder};
use std::{array::TryFromSliceError, collections::HashMap, convert::Infallible};
use thiserror::Error;

#[derive(Error, Debug)]
pub enum RecordEncryptError {
    #[error("EncryptionError: {0}")]
    EncryptionError(#[from] EncryptionError),
    #[error("Failed to encode fields: {0}")]
    FieldEncodeError(String),
    #[error("Failed to encode encrypted record: {0}")]
    EncryptedRecordEncodeError(String),
}

#[derive(Error, Debug)]
pub enum RecordDecryptError {
    #[error("Record id was invalid: {0}")]
    InvalidIdError(String),
    #[error("Failed to decode encrypted record: {0}")]
    DecodeEncryptedRecordError(String),
    #[error("Failed to decrypt record: {0}")]
    DecryptionError(#[from] DecryptionError),
    #[error("Failed to decode record: {0}")]
    DecodeRecordError(String),
}

#[derive(Debug, PartialEq)]
pub struct Record {
    pub id: Uuid,
    pub fields: HashMap<String, Value>,
}

impl Record {
    pub fn get(&self, field: &str) -> Option<&Value> {
        self.fields.get(field)
    }

    /// Decrypt a record from an RPC Document
    pub async fn decrypt(
        cipher: &EnvelopeCipher<impl KeyProvider>,
        doc: Document,
    ) -> Result<Self, RecordDecryptError> {
        let id = Uuid::from_bytes(
            doc.id
                .as_slice()
                .try_into()
                .map_err(|e: TryFromSliceError| {
                    RecordDecryptError::InvalidIdError(e.to_string())
                })?,
        );

        Self::try_from_cbor_bytes(
            id,
            &cipher
                .decrypt(
                    &EncryptedRecord::from_vec(doc.source).map_err(|e| {
                        RecordDecryptError::DecodeEncryptedRecordError(e.to_string())
                    })?,
                )
                .await?,
        )
        .map_err(|e| RecordDecryptError::DecodeRecordError(e.to_string()))
    }

    /// Encrypt a record as an RPC Document
    pub async fn encrypt(
        &self,
        cipher: &EnvelopeCipher<impl KeyProvider>,
    ) -> Result<Document, RecordEncryptError> {
        Ok(Document {
            id: self.id.as_bytes().to_vec(),
            source: cipher
                .encrypt(
                    &self
                        .as_cbor_bytes()
                        .map_err(|e| RecordEncryptError::FieldEncodeError(e.to_string()))?,
                )
                .await?
                .to_vec()
                .map_err(|e| RecordEncryptError::EncryptedRecordEncodeError(e.to_string()))?,
        })
    }

    /// Losslessly encode a record as CBOR
    ///
    /// When the record is encoded using this method it is safe to decode it using the
    /// `try_from_cbor_bytes` method.
    /// This means the bytes returned by this method can be encrypted and stored in the data
    /// service without needing to use the [`RecordDecoder`] upon retrieval.
    pub fn as_cbor_bytes(&self) -> Result<Vec<u8>, minicbor::encode::Error<Infallible>> {
        let mut output = vec![];

        let mut encoder = Encoder::new(&mut output);

        encoder.map(self.fields.len() as _)?;

        for (k, v) in self.fields.iter() {
            encoder.encode(k)?;
            encoder.encode(v)?;
        }

        Ok(output)
    }

    /// Load a record from CBOR without validating against a schema.
    ///
    /// Note: this method should only be used when the CBOR bytes were encoded with the correct types.
    /// When CBOR is coming from an external source (such as JavaScript) the [`RecordDecoder`]
    /// should instead be used.
    pub fn try_from_cbor_bytes(id: Uuid, bytes: &[u8]) -> Result<Self, minicbor::decode::Error> {
        let mut decoder = Decoder::new(bytes);

        Ok(Self {
            id,
            fields: decoder.map_iter()?.collect::<Result<_, _>>()?,
        })
    }

    /// Index into a record using dot notation
    ///
    /// # Examples
    ///
    /// ```ignore
    /// // Record {
    /// //   first: {
    /// //     second: Value::String("hey")
    /// //   }
    /// // }
    ///
    /// record.index_with_dot_notation("first.second") // Some(Value::String("hey"))
    /// ```
    pub(crate) fn index_with_dot_notation<'a>(&'a self, target: &str) -> Option<&'a Value> {
        let parts = target.split('.').collect::<Vec<_>>();

        fn traverse_record<'b>(
            record: &'b HashMap<String, Value>,
            parts: &[&str],
        ) -> Option<&'b Value> {
            if parts.len() == 1 {
                record.get(parts[0])
            } else if let Some((first, tail)) = parts.split_first() {
                if let Some(Value::Map(inner_record)) = record.get(*first) {
                    traverse_record(inner_record, tail)
                } else {
                    None
                }
            } else {
                None
            }
        }

        traverse_record(&self.fields, &parts)
    }

    /// Return the values of the specified fields if they are Value::String.
    pub(crate) fn extract_string_fields(
        &self,
        fields: impl IntoIterator<Item = impl AsRef<str>>,
    ) -> Vec<String> {
        let mut output = vec![];

        for field in fields.into_iter() {
            if let Some(Value::String(s)) = self.index_with_dot_notation(field.as_ref()) {
                output.push(s.clone());
            }
        }

        output
    }

    /// Return all the fields inside the record that are Value::String.
    pub(crate) fn extract_all_string_fields(&self) -> Vec<String> {
        let mut output = vec![];

        fn traverse(val: &Value, output: &mut Vec<String>) {
            match val {
                Value::String(s) => {
                    output.push(s.clone());
                }
                Value::Map(map) => {
                    for val in map.values() {
                        traverse(val, output);
                    }
                }
                _ => {
                    // Not a string field - so safe to ignore.
                }
            }
        }

        for val in self.fields.values() {
            traverse(val, &mut output);
        }

        output
    }

    /// Return all the fields inside the record that are Value::String and their keys.
    ///
    /// For nested fields, the keys of the strings are represented using dot notation.
    pub(crate) fn extract_all_string_fields_and_keys(&self) -> Vec<(String, String)> {
        let mut output = vec![];
        let mut key = vec![];

        fn traverse(key: &mut Vec<String>, val: &Value, output: &mut Vec<(String, String)>) {
            match val {
                Value::String(s) => {
                    let key_string = helpers::join_with_dot(key);
                    let value = s.clone();

                    output.push((key_string, value));
                }
                Value::Map(map) => {
                    for (k, v) in map.iter() {
                        key.push(k.into());

                        traverse(key, v, output);

                        key.pop();
                    }
                }
                _ => {
                    // Not a string field - so safe to ignore.
                }
            }
        }

        for (k, v) in self.fields.iter() {
            key.push(k.into());

            traverse(&mut key, v, &mut output);

            key.pop();
        }

        output
    }
}

#[cfg(test)]
mod tests {
    use super::{Record, RecordDecryptError, Value};
    use crate::utils::{collection, record};
    use cipherstash_grpc::api::documents::Document;
    use envelopers::{EnvelopeCipher, SimpleKeyProvider};
    use uuid::Uuid;

    fn sorted<T: Ord>(mut vec: Vec<T>) -> Vec<T> {
        vec.sort();
        vec
    }

    fn create_nested_record() -> Record {
        record! {
            "id" => Uuid::from_bytes([0; 16]),
            "int" => 10,
            "float" => 10.,
            "bool" => true,
            "string" => "Top level!",
            "date" => Value::date_from_millis(100),
            "nested" => collection! {
                "int" => 10,
                "float" => 10.,
                "bool" => true,
                "string" => "Second level!",
                "another" => collection! {
                    "int" => 10,
                    "float" => 10.,
                    "bool" => true,
                    "string" => "Third level!"
                }
            }
        }
    }

    #[test]
    fn test_index_with_dot_notation() {
        let record = record! {
            "id" => Uuid::from_bytes([0; 16]),
            "first" => 10,
            "nested" => collection! {
                "second" => 10.,
                "another" => collection! {
                    "third" => false
                }
            }
        };

        assert_eq!(
            record.index_with_dot_notation("first"),
            Some(&Value::Uint64(10))
        );

        assert_eq!(
            record.index_with_dot_notation("nested.second"),
            Some(&Value::Float64(10.))
        );

        assert_eq!(
            record.index_with_dot_notation("nested.another.third"),
            Some(&Value::Boolean(false))
        );
    }

    #[test]
    fn test_extract_string_fields_from_keys() {
        let record = create_nested_record();

        assert_eq!(
            sorted(record.extract_string_fields(["string", "nested.string"])),
            sorted(vec![
                String::from("Top level!"),
                String::from("Second level!")
            ])
        );
    }

    #[test]
    fn test_extract_all_string_fields() {
        let record = create_nested_record();

        assert_eq!(
            sorted(record.extract_all_string_fields()),
            sorted(vec![
                String::from("Top level!"),
                String::from("Second level!"),
                String::from("Third level!"),
            ])
        )
    }

    #[test]
    fn test_extract_all_string_fields_with_keys() {
        let record = create_nested_record();

        assert_eq!(
            sorted(record.extract_all_string_fields_and_keys()),
            sorted(vec![
                (String::from("string"), String::from("Top level!")),
                (String::from("nested.string"), String::from("Second level!")),
                (
                    String::from("nested.another.string"),
                    String::from("Third level!")
                ),
            ])
        )
    }

    #[test]
    fn test_round_trip_nested_record() {
        let record = create_nested_record();

        let copy = Record::try_from_cbor_bytes(
            record.id,
            &record
                .as_cbor_bytes()
                .expect("Failed to encode record as CBOR"),
        )
        .expect("Failed to load record from CBOR");

        assert_eq!(record, copy);
    }

    #[tokio::test]
    async fn test_round_trip_nested_record_encryption() {
        let record = create_nested_record();

        let cipher: EnvelopeCipher<SimpleKeyProvider> =
            EnvelopeCipher::init(SimpleKeyProvider::init([2; 16]));

        let copy = Record::decrypt(
            &cipher,
            record
                .encrypt(&cipher)
                .await
                .expect("Failed to encrypt record"),
        )
        .await
        .expect("Failed to decrypt record");

        assert_eq!(record, copy);
    }

    #[tokio::test]
    async fn test_decrypt_invalid_key() {
        let record = create_nested_record();

        let first_cipher: EnvelopeCipher<SimpleKeyProvider> =
            EnvelopeCipher::init(SimpleKeyProvider::init([2; 16]));

        let second_cipher: EnvelopeCipher<SimpleKeyProvider> =
            EnvelopeCipher::init(SimpleKeyProvider::init([3; 16]));

        let encrypted = record
            .encrypt(&first_cipher)
            .await
            .expect("Failed to encrypt record");

        let err = Record::decrypt(&second_cipher, encrypted)
            .await
            .expect_err("Expected decrypt to fail");

        assert_eq!(
            err.to_string(),
            "Failed to decrypt record: failed to decrypt data key"
        )
    }

    #[tokio::test]
    async fn test_decrypt_invalid_encrypted_record() {
        let cipher: EnvelopeCipher<SimpleKeyProvider> =
            EnvelopeCipher::init(SimpleKeyProvider::init([2; 16]));

        let err = Record::decrypt(
            &cipher,
            Document {
                id: vec![1; 16],
                source: vec![1, 2, 3],
            },
        )
        .await
        .expect_err("Expected decrypt to fail");

        assert!(matches!(
            err,
            RecordDecryptError::DecodeEncryptedRecordError(_)
        ))
    }

    #[tokio::test]
    async fn test_decrypt_invalid_cbor() {
        let cipher: EnvelopeCipher<SimpleKeyProvider> =
            EnvelopeCipher::init(SimpleKeyProvider::init([2; 16]));

        let err = Record::decrypt(
            &cipher,
            Document {
                id: vec![1; 16],
                source: cipher
                    .encrypt(&[1, 2, 3, 4])
                    .await
                    .expect("Failed to encrypt")
                    .to_vec()
                    .unwrap(),
            },
        )
        .await
        .expect_err("Expected decrypt to fail");

        assert_eq!(
            err.to_string(),
            "Failed to decode record: unexpected type u8 at position 0: expected map"
        )
    }
}