//! Contains structs for indexing records based on a collection& schema.

pub mod errors;
pub(crate) mod fixed_byte_buf;
pub mod mapping;
pub mod mapping_indexer;
pub mod query;
pub mod text;
pub mod vector;

use crate::record::decoder::SchemaField;

use super::record::{
    decoder::{RecordDecoder, RecordSchema},
    Record,
};
use errors::{CollectionSchemaError, LoadIndexerError};
use mapping::MappingWithMeta;
use mapping_indexer::MappingIndexer;

use ciborium::{de::from_reader, ser::into_writer};
use serde::Deserialize;
use std::collections::{hash_map::RandomState, HashMap};

use errors::{QueryIndexingError, RecordIndexingError};
use query::{ConjunctiveCondition, Query};
use vector::{Constraint, TermVector};

use hmac::{Hmac, Mac};
use sha2::Sha256;
type HmacSha256 = Hmac<Sha256>;

/// The schema for a collection - combining the schema for individual records with the collection's
/// index specifications.
#[derive(Deserialize, Debug)]
pub struct CollectionSchema {
    // FIXME: Made these public for the Cloudflare demo (discuss)
    pub indexes: HashMap<String, MappingWithMeta>,
    #[serde(rename = "type")]
    pub schema: RecordSchema,
}

impl CollectionSchema {
    pub fn with_index_key(
        key: [u8; 32],
        mut indexes: HashMap<String, MappingWithMeta>,
        schema: HashMap<String, SchemaField, RandomState>,
    ) -> Result<Self, CollectionSchemaError> {
        for (name, meta) in indexes.iter_mut() {
            let mut mac = HmacSha256::new_from_slice(&key)
                .map_err(|_| CollectionSchemaError::InvalidIndexKey)?;
            mac.update(format!("{}:ORE", name).as_bytes());
            let field_key = mac.finalize().into_bytes();
            meta.prf_key.copy_from_slice(&field_key[0..16]);
            meta.prp_key.copy_from_slice(&field_key[16..]);
        }

        Ok(Self {
            indexes,
            schema: RecordSchema { map: schema },
        })
    }
}

/// An indexer that can index an entire record based on a [`CollectionSchema`]
pub struct Indexer {
    indexes: HashMap<String, MappingIndexer>,
    decoder: RecordDecoder,
}

impl Indexer {
    /// Create an indexer from a [`CollectionSchema`]
    pub fn from_collection_schema(
        CollectionSchema { indexes, schema }: CollectionSchema,
    ) -> Result<Self, LoadIndexerError> {
        let decoder = RecordDecoder { schema };

        let indexes = indexes
            .into_iter()
            .map(|(key, mapping)| {
                MappingIndexer::from_mapping(mapping).map(|indexer| (key, indexer))
            })
            .collect::<Result<HashMap<_, _>, _>>()?;

        Ok(Self { decoder, indexes })
    }

    /// Decode a record indexer from a CBOR byte array
    ///
    /// The CBOR byte array is a serialized [`CollectionSchema`] struct.
    pub fn decode_from_cbor(bytes: &[u8]) -> Result<Self, LoadIndexerError> {
        let schema: CollectionSchema = from_reader(bytes)?;

        Self::from_collection_schema(schema)
    }

    /// Encrypt a record and return an array of [`TermVector`] used for indexing
    pub fn encrypt(&mut self, record: &Record) -> Result<Vec<TermVector>, RecordIndexingError> {
        self.indexes
            .values_mut()
            .filter_map(|index| index.encrypt(record).transpose())
            .collect()
    }

    /// Encrypt a query and return an array of [`Constraint`] used for querying
    pub fn encrypt_query(&mut self, query: Query) -> Result<Vec<Constraint>, QueryIndexingError> {
        let mut output: Vec<Constraint> = vec![];

        fn traverse_query(
            query: Query,
            indexes: &mut HashMap<String, MappingIndexer>,
            output: &mut Vec<Constraint>,
        ) -> Result<(), QueryIndexingError> {
            match query {
                Query::Basic { index_name, kind } => {
                    let mapping = indexes
                        .get_mut(&index_name)
                        .ok_or(QueryIndexingError::InvalidIndexName(index_name))?;

                    if let Some(c) = mapping.encrypt_query(kind)? {
                        output.extend(c.into_iter());
                    }
                }
                Query::Conjunctive(ConjunctiveCondition::All { conditions }) => {
                    for q in conditions.into_iter() {
                        traverse_query(q, indexes, output)?;
                    }
                }
            }

            Ok(())
        }

        traverse_query(query, &mut self.indexes, &mut output)?;

        Ok(output)
    }

    /// Decode a record from CBOR, encrypt it using ORE and return the result as a CBOR byte
    /// buffer.
    ///
    /// The CBOR byte array is a serialized array of [`TermVector`]s.
    pub fn encrypt_cbor(&mut self, record_bytes: &[u8]) -> Result<Vec<u8>, RecordIndexingError> {
        let mut output = vec![];
        let record = self.decoder.decode(record_bytes)?;
        into_writer(&self.encrypt(&record)?, &mut output)?;
        Ok(output)
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        record::{
            decoder::{DataType, RecordSchema, SchemaField},
            Record,
        },
        test_utils::{cbor_buffer, collection},
    };

    use super::{
        errors::QueryIndexingError,
        mapping::{Mapping, MappingWithMeta},
        query::*,
        CollectionSchema, Indexer,
    };

    use serde_bytes::Bytes;
    use uuid::uuid;

    #[test]
    fn test_load_schema_from_cbor() {
        let buf = Bytes::new(&[0_u8; 16]);

        let buffer = cbor_buffer!({
            "type" => {
                "title" => "string",
                "runningTime" => "uint64",
                "year" => "uint64"
            },
            "indexes" => {
                "exactTitle" => {
                    "kind" => "exact",
                    "field" => "title",
                    "prf_key" => &buf,
                    "prp_key" => &buf,
                    "index_id" => &buf
                },
                "runningTime" => {
                    "kind" => "range",
                    "field" => "runningTime",
                    "prf_key" => &buf,
                    "prp_key" => &buf,
                    "index_id" => &buf
                }
            }
        });

        let indexer = Indexer::decode_from_cbor(&buffer).expect("Failed to decode");

        assert_eq!(indexer.indexes.len(), 2);

        let exact = indexer
            .indexes
            .get("exactTitle")
            .expect("Expected exactTitle index to exist");

        assert_eq!(
            exact.mapping,
            Mapping::Exact {
                field: String::from("title")
            }
        );

        assert_eq!(exact.index_id, [0; 16]);

        let time = indexer
            .indexes
            .get("runningTime")
            .expect("Expected runningTime index to exist");

        assert_eq!(
            time.mapping,
            Mapping::Range {
                field: String::from("runningTime")
            }
        );
    }

    #[test]
    fn test_load_schema_from_json() {
        let value = serde_json::json!({
            "type": {
                "title": "string",
                "runningTime": "uint64",
                "year": "uint64"
            },
            "indexes": {
                "exactTitle": {
                    "kind": "exact",
                    "field": "title",
                    "$prfKey": [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ],
                    "$prpKey": [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ],
                    "$indexId": "00000000-0000-0000-0000-000000000000"
                }
            }
        });

        let schema: CollectionSchema =
            serde_json::from_value(value).expect("Failed to decode schema");

        let indexer = Indexer::from_collection_schema(schema).expect("Failed to load indexer");

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

        let exact = indexer
            .indexes
            .get("exactTitle")
            .expect("Expected exactTitle index to exist");

        assert_eq!(exact.index_id, [0; 16]);
    }

    fn create_test_movie_indexer() -> Indexer {
        let collection_schema = CollectionSchema {
            schema: RecordSchema {
                map: collection! {
                    "title" => SchemaField::DataType(DataType::String),
                    "runningTime" => SchemaField::DataType(DataType::Uint64)
                },
            },
            indexes: collection! {
                "exactTitle" => MappingWithMeta {
                    mapping: Mapping::Exact { field: "title".into() },
                    index_id: uuid!("00000000-0000-0000-0000-000000000000"),
                    prf_key: [0;16],
                    prp_key: [0;16]
                },
                "runningTime" => MappingWithMeta {
                    mapping: Mapping::Range { field: "runningTime".into() },
                    index_id: uuid!("00000000-0000-0000-0000-000000000000"),
                    prf_key: [0;16],
                    prp_key: [0;16]
                }
            },
        };

        Indexer::from_collection_schema(collection_schema).unwrap()
    }

    fn create_test_user_indexer() -> Indexer {
        let collection_schema = CollectionSchema {
            schema: RecordSchema {
                map: collection! {
                    "firstName" => SchemaField::DataType(DataType::String),
                    "lastName" => SchemaField::DataType(DataType::String),
                    "dob" => SchemaField::DataType(DataType::Date),
                    "stashId" => SchemaField::DataType(DataType::String)
                },
            },
            indexes: collection! {
                "firstName" => MappingWithMeta {
                    mapping: Mapping::Exact { field: "firstName".into() },
                    index_id: uuid!("00000000-0000-0000-0000-000000000000"),
                    prf_key: [0;16],
                    prp_key: [0;16]
                }
            },
        };

        Indexer::from_collection_schema(collection_schema).unwrap()
    }

    #[test]
    fn index_record_with_both_indexes() {
        let mut indexer = create_test_movie_indexer();

        let vectors = indexer
            .encrypt(&Record {
                id: [
                    163, 98, 105, 100, 216, 64, 80, 117, 189, 53, 179, 82, 84, 65, 2, 178,
                ],
                fields: collection! {
                    "title" => "Hello!",
                    "runningTime" => 230
                },
            })
            .unwrap();

        assert_eq!(vectors.len(), 2);
    }

    #[test]
    fn index_user_record_from_cbor() {
        let mut indexer = create_test_user_indexer();

        let cbuf = cbor_buffer!({
            "id" => [163,  98, 105, 100, 216, 64,
                    80, 117, 189,  53, 179, 82,
                    84,  65,   2, 178],
            "firstName" => "Dan",
            "lastName" => "Draper",
            "stashId" => "e976a7eb-5a17-46bf-aef9-d70fe77d79b5"
        });

        indexer.encrypt_cbor(&cbuf).unwrap();

        // If we get this far then nothing failed
        // TODO: do the cbor decoding and check the result
        assert!(true)
    }

    #[test]
    fn index_record_with_missing_field() {
        let mut indexer = create_test_movie_indexer();

        let vectors = indexer
            .encrypt(&Record {
                id: [1; 16],
                fields: collection! {
                    "title" => "Hello!"
                },
            })
            .unwrap();

        assert_eq!(vectors.len(), 1);
    }

    #[test]
    fn index_record_with_no_fields() {
        let mut indexer = create_test_movie_indexer();

        let vectors = indexer
            .encrypt(&Record {
                id: [1; 16],
                fields: collection! {},
            })
            .unwrap();

        assert_eq!(vectors.len(), 0);
    }

    #[test]
    fn test_query_collapse_all_query() {
        let mut indexer = create_test_movie_indexer();

        let constraints = indexer
            .encrypt_query(Query::Conjunctive(ConjunctiveCondition::All {
                conditions: vec![
                    Query::Basic {
                        index_name: "exactTitle".into(),
                        kind: QueryKind::Exact {
                            value: "test".into(),
                        },
                    },
                    Query::Basic {
                        index_name: "runningTime".into(),
                        kind: QueryKind::Range {
                            value: RangeValue::Between {
                                min: 0.into(),
                                max: 100.into(),
                            },
                        },
                    },
                ],
            }))
            .expect("Failed to encrypt query");

        assert_eq!(constraints.len(), 2)
    }

    #[test]
    fn test_nested_all_are_collapsed() {
        let mut indexer = create_test_movie_indexer();

        let constraints = indexer
            .encrypt_query(Query::Conjunctive(ConjunctiveCondition::All {
                conditions: vec![Query::Conjunctive(ConjunctiveCondition::All {
                    conditions: vec![Query::Conjunctive(ConjunctiveCondition::All {
                        conditions: vec![Query::Basic {
                            index_name: "exactTitle".into(),
                            kind: QueryKind::Exact {
                                value: "test".into(),
                            },
                        }],
                    })],
                })],
            }))
            .expect("Failed to encrypt query");

        assert_eq!(constraints.len(), 1)
    }

    #[test]
    fn test_encrypt_single_query() {
        let mut indexer = create_test_movie_indexer();

        let constraints = indexer
            .encrypt_query(Query::Basic {
                index_name: "exactTitle".into(),
                kind: QueryKind::Exact {
                    value: "test".into(),
                },
            })
            .expect("Failed to encrypt query");

        assert_eq!(constraints.len(), 1)
    }

    #[test]
    fn test_fail_when_querying_invalid_index() {
        let mut indexer = create_test_movie_indexer();

        let err = indexer
            .encrypt_query(Query::Basic {
                index_name: "not-a-real-index".into(),
                kind: QueryKind::Exact {
                    value: "test".into(),
                },
            })
            .expect_err("Expected encryption to fail");

        assert_eq!(
            err,
            QueryIndexingError::InvalidIndexName("not-a-real-index".into())
        );
    }

    #[test]
    fn test_fail_when_index_is_wrong_type() {
        let mut indexer = create_test_movie_indexer();

        let err = indexer
            .encrypt_query(Query::Basic {
                index_name: "exactTitle".into(),
                kind: QueryKind::Range {
                    value: RangeValue::Single {
                        operator: RangeOperator::Eq,
                        value: 10.into(),
                    },
                },
            })
            .expect_err("Expected encryption to fail");

        assert_eq!(
            err,
            QueryIndexingError::InvalidQueryMappingPair {
                query_name: "range",
                mapping_name: "exact"
            }
        );
    }

    fn collection_schema_from_index_key() {
        let index_key = [0u8; 32];
        let collection_schema = CollectionSchema::with_index_key(
            index_key,
            collection! {
                "exactTitle" => MappingWithMeta {
                    mapping: Mapping::Exact { field: "title".into() },
                    index_id: uuid!("00000000-0000-0000-0000-000000000000"),
                    prf_key: [0;16],
                    prp_key: [0;16]
                },
                "runningTime" => MappingWithMeta {
                    mapping: Mapping::Range { field: "runningTime".into() },
                    index_id: uuid!("00000000-0000-0000-0000-000000000000"),
                    prf_key: [0;16],
                    prp_key: [0;16]
                }
            },
            collection! {
                "title" => SchemaField::DataType(DataType::String),
                "runningTime" => SchemaField::DataType(DataType::Uint64)
            },
        )
        .unwrap();

        if let Some(meta) = collection_schema.indexes.get("exactTitle") {
            // TODO: Handle the index_id too
            assert_ne!(meta.prf_key, [0u8; 16]);
            assert_ne!(meta.prp_key, [0u8; 16]);
        } else {
            assert!(false);
        }
    }
}