//! Contains structs for indexing a record based on a single [`Mapping`]

use super::{
    errors::{IndexerError, LoadIndexerError},
    mapping::{Mapping, MappingWithMeta},
    vector::{EncryptedTerm, TermVector, VectorTerm},
};
use crate::record::{Record, Value};

use hex_literal::hex;
use lazy_static::lazy_static;
use num_bigint::{BigUint, ToBigUint};
use ore_encoding_rs::{siphash, OrePlaintext};
use ore_rs::{scheme::bit2::OREAES128, ORECipher, OREEncrypt};
use regex::Regex;

/// Encrypt string terms using siphash and ORE for match queries
fn encrypt_string_terms(
    ore: &mut OREAES128,
    index_id: &[u8; 16],
    record_id: &[u8; 16],
    terms: Vec<String>,
) -> Result<Option<TermVector>, IndexerError> {
    let terms = terms
        .into_iter()
        .map(|term| OrePlaintext::from(siphash(term.as_bytes())))
        .map(|plaintext| {
            plaintext
                .0
                .encrypt(ore)
                .map(EncryptedTerm::from)
                .map(|term| VectorTerm {
                    term,
                    link: *record_id,
                })
        })
        .collect::<Result<Vec<_>, _>>()?;

    // If there are no string terms there is no point inserting anything into the database - just
    // return None.
    if terms.is_empty() {
        return Ok(None);
    }

    Ok(Some(TermVector {
        terms,
        index_id: *index_id,
    }))
}

/// Encrypt u64 terms using ORE for ordered strings
fn encrypt_u64_terms(
    ore: &mut OREAES128,
    index_id: &[u8; 16],
    record_id: &[u8; 16],
    terms: Vec<u64>,
) -> Result<Option<TermVector>, IndexerError> {
    let terms = terms
        .into_iter()
        .map(|term| OrePlaintext::from(term))
        .map(|plaintext| plaintext.0.encrypt(ore))
        .collect::<Result<Vec<_>, _>>()?;

    // If there are no terms there is no point inserting anything into the database - just
    // return None.
    if terms.is_empty() {
        return Ok(None);
    }

    Ok(Some(TermVector {
        terms: vec![VectorTerm {
            term: EncryptedTerm(terms),
            link: *record_id,
        }],
        index_id: *index_id,
    }))
}

fn orderise_string(s: &str) -> Result<Vec<u64>, IndexerError> {
    if !s.is_ascii() {
        return Err(IndexerError::InvalidRecordError(
            "Can only order strings that are pure ASCII".to_string(),
        ));
    }

    lazy_static! {
        static ref NON_ALPHANUMERIC_OR_SPACE_REGEX: Regex = Regex::new("[^a-z0-9[:space:]]+").unwrap();
    }

    lazy_static! {
        static ref SPACE_REGEX: Regex = Regex::new("[[:space:]]+").unwrap();
    }

    lazy_static! {
        static ref DIGIT_REGEX: Regex = Regex::new("[0-9]").unwrap();
    }

    // This all very much relies on ASCII character numbering.  A copy of `ascii`(7)
    // up on a convenient terminal may assist in understanding what's going
    // on here.

    // First up, let's transmogrify the string we were given into one that only contains
    // a controlled subset of characters, that we can easily map into a smaller numeric
    // space.

    let mut s = s.to_lowercase();

    // Any group of rando characters sort at the end.
    s = NON_ALPHANUMERIC_OR_SPACE_REGEX
        .replace_all(&s, "~")
        .to_string();

    // Any amount of whitespace comes immediately after letters.
    s = SPACE_REGEX.replace_all(&s, "{").to_string();

    // Numbers come after spaces.
    s = DIGIT_REGEX.replace_all(&s, "|").to_string();

    // Next, we turn that string of characters into a "packed" number that represents the
    // whole string, but in a more compact form than would be used if each character took
    // up the full seven or eight bits used by regular ASCII.
    let mut n = s
        .bytes()
        // 'a' => 1, 'b' => 2, ..., 'z' => 27, '{' => 28, '|' => 29,
        // '}' => 30 (unused), '~' => 31.  0 is kept as "no character" so
        // that short strings sort before longer ones.
        .map(|c| BigUint::from(c - 96))
        // Turn the whole thing into one giant number, with each character
        // occupying five bits of said number.
        .fold(0.to_biguint().unwrap(), |i, c| (i << 5) + c);

    // Thirdly, we need to turn the number into one whose in-memory representation
    // has a length in bits that is a multiple of 64.  This is to ensure that
    // the first character has the most-significant bits possible, so it
    // sorts the highest.
    n = n << (64 - (s.len() * 5) % 64);

    let mut terms = Vec::new();

    let two_pow_64 = 2.to_biguint().unwrap().pow(64);

    // And now, semi-finally, we can turn all that gigantic mess into a vec of u64 terms.
    while n > 0.to_biguint().unwrap() {
        // Unwrapping is fine here because we'll end up with a number that
        // fits into a u64.
        let term = (&n % &two_pow_64).try_into().unwrap();

        terms.insert(0, term);

        n >>= 64;
    }

    // Only six ORE ciphertexts can fit into the database.
    terms.truncate(6);

    Ok(terms)
}

/// An Indexer that indexs a record based on a single index.
pub struct MappingIndexer {
    ore: OREAES128,
    pub(crate) index_id: [u8; 16],
    pub(crate) mapping: Mapping,
}

impl MappingIndexer {
    /// Load a single MappingIndexer from a [`MappingWithMeta`]
    pub fn from_mapping(mapping: MappingWithMeta) -> Result<Self, LoadIndexerError> {
        let seed = hex!("00010203 04050607");

        let ore = OREAES128::init(mapping.prf_key, mapping.prp_key, &seed)
            .map_err(|e| LoadIndexerError::Other(format!("ORE init failed: {:?}", e)))?;

        Ok(Self {
            ore,
            index_id: mapping.index_id,
            mapping: mapping.mapping,
        })
    }

    /// Create a [`TermVector`] from a single term for the current mapping
    pub fn vector_from_term(&self, term: VectorTerm) -> TermVector {
        TermVector {
            terms: vec![term],
            index_id: self.index_id,
        }
    }

    /// Encrypt a [`Record`] using ORE
    ///
    /// If the Record contains no meaningful data to encrypt for this index, this method will
    /// return `None`.
    pub fn encrypt(&mut self, record: &Record) -> Result<Option<TermVector>, IndexerError> {
        let id = record.id;

        Ok(match &self.mapping {
            Mapping::Exact { field } => {
                if let Some(plaintext) = record
                    .index_with_dot_notation(field)
                    .and_then(|x| x.as_plaintext())
                {
                    let term = plaintext.0.encrypt(&mut self.ore)?.into();

                    Some(self.vector_from_term(VectorTerm { term, link: id }))
                } else {
                    None
                }
            }

            Mapping::Range { field } => {
                if let Some(val) = record.index_with_dot_notation(field) {
                    if let Value::String(s) = val {
                        let terms = orderise_string(s)?;
                        encrypt_u64_terms(&mut self.ore, &self.index_id, &record.id, terms)?
                    } else if let Some(plaintext) = val.as_plaintext() {
                        let term = plaintext.0.encrypt(&mut self.ore)?.into();
                        Some(self.vector_from_term(VectorTerm { term, link: id }))
                    } else {
                        // If the current value doesn't have a plaintext representation (it may be
                        // a map or array) then just exclude it from the index.
                        None
                    }
                } else {
                    None
                }
            }

            Mapping::Match { fields, pipeline } => encrypt_string_terms(
                &mut self.ore,
                &self.index_id,
                &record.id,
                // Match only operates on specified fields and only if they are strings
                pipeline.process(record.extract_string_fields(fields)),
            )?,

            Mapping::DynamicMatch { pipeline } => encrypt_string_terms(
                &mut self.ore,
                &self.index_id,
                &record.id,
                // Dynamic match operates on all string fields
                pipeline.process(record.extract_all_string_fields()),
            )?,

            Mapping::FieldDynamicMatch { pipeline } => encrypt_string_terms(
                &mut self.ore,
                &self.index_id,
                &record.id,
                record
                    .extract_all_string_fields_and_keys()
                    .into_iter()
                    .flat_map(|(k, v)| {
                        pipeline
                            .process(vec![v])
                            .into_iter()
                            // In order to specifiy fields and terms to be searches they must be
                            // encrypted in this "field:value" format. If this isn't done querying will
                            // silently fail.
                            .map(move |t| format!("{}:{}", k, t))
                    })
                    .collect(),
            )?,
        })
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        indexer::{
            errors,
            mapping::{Mapping, MappingWithMeta},
        },
        record::Record,
        test_utils::collection,
    };

    use super::{orderise_string, MappingIndexer};
    use serde::Deserialize;
    use std::fs;
    use std::path::Path;

    #[test]
    fn test_compare_exact_matches() {
        let record = Record {
            id: [0; 16],
            fields: collection! {
                "test" => "test-string"
            },
        };

        let mut left_indexer = MappingIndexer::from_mapping(MappingWithMeta {
            mapping: Mapping::Exact {
                field: "test".into(),
            },
            prp_key: [0; 16],
            prf_key: [1; 16],
            index_id: [2; 16],
        })
        .expect("Failed to gen left indexer");

        let mut right_indexer = MappingIndexer::from_mapping(MappingWithMeta {
            mapping: Mapping::Exact {
                field: "test".into(),
            },
            prp_key: [0; 16],
            prf_key: [1; 16],
            index_id: [2; 16],
        })
        .expect("Failed to gen right indexer");

        let left = left_indexer
            .encrypt(&record)
            .expect("Failed to encrypt")
            .unwrap();
        let right = right_indexer
            .encrypt(&record)
            .expect("Failed to encrypt")
            .unwrap();

        assert_eq!(left.terms[0], right.terms[0]);
    }

    #[test]
    fn test_compare_not_exact() {
        let first = Record {
            id: [0; 16],
            fields: collection! {
                "test" => "test-string"
            },
        };

        let second = Record {
            id: [1; 16],
            fields: collection! {
                "test" => "test-different-string"
            },
        };

        let mut indexer = MappingIndexer::from_mapping(MappingWithMeta {
            mapping: Mapping::Exact {
                field: "test".into(),
            },
            prp_key: [0; 16],
            prf_key: [1; 16],
            index_id: [2; 16],
        })
        .expect("Failed to gen left indexer");

        let left = indexer.encrypt(&first).expect("Failed to encrypt").unwrap();
        let right = indexer
            .encrypt(&second)
            .expect("Failed to encrypt")
            .unwrap();

        assert_ne!(left.terms[0], right.terms[0]);
    }

    #[test]
    fn test_compare_range() {
        let first = Record {
            id: [0; 16],
            fields: collection! {
                "test" => 10
            },
        };

        let second = Record {
            id: [1; 16],
            fields: collection! {
                "test" => 20
            },
        };

        let mut indexer = MappingIndexer::from_mapping(MappingWithMeta {
            mapping: Mapping::Range {
                field: "test".into(),
            },
            prp_key: [0; 16],
            prf_key: [1; 16],
            index_id: [2; 16],
        })
        .expect("Failed to gen left indexer");

        let left = indexer.encrypt(&first).expect("Failed to encrypt").unwrap();
        let right = indexer
            .encrypt(&second)
            .expect("Failed to encrypt")
            .unwrap();

        assert!(left.terms[0] < right.terms[0]);
    }

    #[test]
    fn test_compare_range_with_strings() {
        let second = Record {
            id: [0; 16],
            fields: collection! {
                "test" => "b"
            },
        };

        let first = Record {
            id: [1; 16],
            fields: collection! {
                "test" => "a"
            },
        };

        let third = Record {
            id: [2; 16],
            fields: collection! {
                "test" => "c"
            },
        };

        let mut indexer = MappingIndexer::from_mapping(MappingWithMeta {
            mapping: Mapping::Range {
                field: "test".into(),
            },
            prp_key: [0; 16],
            prf_key: [1; 16],
            index_id: [2; 16],
        })
        .expect("Failed to gen left indexer");

        let first_encrypted = indexer.encrypt(&first).expect("Failed to encrypt").unwrap();
        let second_encrypted = indexer
            .encrypt(&second)
            .expect("Failed to encrypt")
            .unwrap();
        let third_encrypted = indexer.encrypt(&third).expect("Failed to encrypt").unwrap();

        assert!(first_encrypted.terms[0].term < second_encrypted.terms[0].term);
        assert!(second_encrypted.terms[0].term < third_encrypted.terms[0].term);
    }

    #[test]
    // Asserts that two ORE encrypted strings from range index compare the same as the ruby client.
    // Tests against a file that contains a snapshot of results from the ruby client.
    fn test_range_comparison_gets_same_result_as_ruby_client() {
        #[derive(Deserialize, Debug)]
        struct TestCase {
            input: Vec<String>,
            output: String,
        }

        let case_file_path =
            Path::new(env!("CARGO_MANIFEST_DIR")).join("./string_comparison_test_cases.json");
        let json_str = fs::read_to_string(case_file_path).expect("couldn't read test case file");
        let test_cases: Vec<TestCase> =
            serde_json::from_str(&json_str).expect("couldn't parse test cases");

        for test_case in test_cases {
            let mut indexer = MappingIndexer::from_mapping(MappingWithMeta {
                mapping: Mapping::Range {
                    field: "test".into(),
                },
                prp_key: [0; 16],
                prf_key: [1; 16],
                index_id: [2; 16],
            })
            .expect("Failed to gen left indexer");

            let str_a = &test_case.input[0];
            let str_b = &test_case.input[1];

            let record_a = Record {
                id: [0; 16],
                fields: collection! {
                    "test" => str_a.clone()
                },
            };

            let record_b = Record {
                id: [1; 16],
                fields: collection! {
                    "test" => str_b.clone()
                },
            };

            let a_encrypted = indexer
                .encrypt(&record_a)
                .expect("Failed to encrypt")
                .unwrap();
            let b_encrypted = indexer
                .encrypt(&record_b)
                .expect("Failed to encrypt")
                .unwrap();

            let a_term = &a_encrypted.terms[0].term;
            let b_term = &b_encrypted.terms[0].term;

            match test_case.output.as_str() {
                "<" => {
                    assert!(
                        a_term < b_term,
                        "expected {:?} to be < {:?} but it was not",
                        &str_a,
                        &str_b,
                    )
                }
                "==" => {
                    assert!(
                        a_term == b_term,
                        "expected {:?} to be == {:?} but it was not",
                        &str_a,
                        &str_b,
                    )
                }
                ">" => {
                    assert!(
                        a_term > b_term,
                        "expected {:?} to be > {:?} but it was not",
                        &str_a,
                        &str_b,
                    )
                }
                op => panic!("unexpected operator: {:?}", op),
            }
        }
    }

    #[test]
    fn test_orderise_string_non_ascii() {
        let result = orderise_string("Jalapeño");

        assert!(matches!(
            result,
            Err(errors::IndexerError::InvalidRecordError(_))
        ));

        let message = result.err().unwrap().to_string();

        assert_eq!(message, "Can only order strings that are pure ASCII")
    }

    #[test]
    // Asserts that output from orderise_string matches output from the ruby client.
    // Tests against a file that contains a snapshot of results from the ruby client.
    fn test_orderise_string_gives_same_output_as_ruby_clint() {
        #[derive(Deserialize, Debug)]
        struct TestCase {
            input: String,
            output: Vec<u64>,
        }

        let case_file_path =
            Path::new(env!("CARGO_MANIFEST_DIR")).join("./orderise_string_test_cases.json");
        let json_str = fs::read_to_string(case_file_path).expect("couldn't read test case file");
        let test_cases: Vec<TestCase> =
            serde_json::from_str(&json_str).expect("couldn't parse test cases");

        for test_case in test_cases {
            let result = orderise_string(&test_case.input);

            assert!(
                result.is_ok(),
                "Expected orderise_string to succeed given {:?}, but got error: {:?}",
                &test_case.input,
                result
            );

            assert_eq!(
                result.unwrap(),
                test_case.output,
                "\n orderise_string didn't match for input: {:?}",
                test_case.input
            );
        }
    }
}