perfume 0.3.2

use async_generic::async_generic;
use base16ct::lower::encode as base16_encode;

use crate::hex_string::HexString;
use crate::random::randomized;
use crate::{Error, STORAGE_KEY_LENGTH};

use super::Identity;
use super::storage::{Storage, StorageState};

// NOTE: implemented with external types to enable codegen before running unit tests. see codegen.rs
/// Compiled data used for random name generation. See [`crate::codegen::ingredients`].
pub type Ingredients = (
    usize,
    phf::Map<&'static str, &'static str>,
    &'static [&'static str],
    &'static [&'static str],
);

/// Persistent random name generator.
pub struct Population<'dom> {
    /// A unique identifier, needed for associating identities with populations.
    pub domain: &'dom str,
    /// Used to generate a keyed hash function, and to randomize word selection.
    pub secret: &'dom [u8],
    /// Words to use for generating names. Created at compile-time with [`crate::codegen::ingredients`].
    pub ingredients: &'static Ingredients,
}

impl<'dom> Population<'dom> {
    /// Generate a unique friendly name from `identifier` which has been persisted using `state`.
    #[async_generic]
    #[allow(unused_assignments)]
    pub fn identity(
        &self,
        identifier: &str,
        state: &mut impl StorageState,
    ) -> Result<Identity<'_>, Error> {
        let storage = self.storage_object(identifier);

        let mut offset = 0usize;
        if _async {
            offset = state.digest_offset_async(self.domain, &storage).await?;
        } else {
            offset = state.digest_offset(self.domain, &storage)?;
        }

        let friendly_name = self.friendly_name(&storage, offset);

        Ok(Identity {
            domain: self.domain,
            friendly_name,
            storage,
        })
    }

    fn storage_object(&self, identifier: &str) -> Storage {
        let mut hasher = blake3::Hasher::new_keyed(self.secret[..32].try_into().unwrap());
        hasher.update(identifier.as_bytes());
        let output = hasher.finalize();
        let mut buf = [0; 64];
        let bytes = base16_encode(output.as_bytes(), &mut buf).unwrap();
        Storage::from(bytes)
    }

    fn friendly_name(&self, storage: &Storage, digest_offset: usize) -> String {
        let (_population_size, prefixes, _colors, _animals) = self.ingredients;

        // prefix comes from a compiled PHF of storage.key -> gerund
        // randomness is provided by the hash function that was used to derive the storage key
        let prefix = prefixes.get(storage.key.as_str()).cloned().unwrap();

        // color and animal are randomly generated by using the storage key and population secret
        // to generate a random u64 value, which is used to select from a compiled list of words
        let animals = self.color_animals(storage);
        let (color, animal) = animals.get(digest_offset).unwrap();

        format!("{prefix}-{color}-{animal}")
    }

    fn color_animals(&self, storage: &Storage) -> Vec<(&str, &str)> {
        let (population_size, _prefixes, colors, animals) = self.ingredients;

        let required_color_animals = *population_size as u32 / 16u32.pow(STORAGE_KEY_LENGTH as u32);

        // use all of the few available colors
        let colors = self.randomize(colors, storage, false);

        // ensure that animals are evenly distributed over colors
        // by using only enough animals to fill a color.
        // NOTE: this implies that the population size can only be chosen once
        let animals_per_color = required_color_animals.div_ceil(colors.len() as u32);
        let animals = self
            .randomize(animals, storage, true)
            .into_iter()
            .take(animals_per_color as usize)
            .collect::<Vec<_>>();

        // fill each color with all available animals before using the next color
        let mut results = vec![];
        for color in colors {
            for &animal in &animals {
                results.push((color, animal))
            }
        }
        results
    }

    fn randomize<'a>(&self, words: &'a [&str], storage: &Storage, reverse: bool) -> Vec<&'a str> {
        // randomization is idempotent because random number seed is based on population "secret"

        // randomized between populations
        let mut buf = [0; 64];
        let pop_seed = base16_encode(&self.secret[..32], &mut buf).unwrap();
        let pop_seed: u16 = HexString::<4>::from(&pop_seed[..4]).into();

        // randomized between storage blobs
        let store_seed = match STORAGE_KEY_LENGTH {
            3 => {
                let mut key_bytes = storage.key.as_str().as_bytes().to_vec(); // 3 bytes
                key_bytes.push(b"0"[0]); // 1 more needed for conversion to u16
                key_bytes
            }
            _ => unimplemented!(),
        };
        let store_seed: u16 = HexString::<4>::from(&store_seed[..4]).into();

        let rng_seed = ((pop_seed as u32) << 16) + (store_seed as u32);
        let mut rng_seed = ((rng_seed as u64) << 32) + (rng_seed as u64);

        // randomized between colors and animals
        if reverse {
            rng_seed = rng_seed.reverse_bits();
        }

        randomized(words, rng_seed)
    }
}

#[cfg(test)]
mod tests {
    /*
    cargo run -F codegen
    cargo test population -- --nocapture
    */

    use std::time::Instant;

    use super::*;
    use crate::identity::{storage::RemoteStore, tests::*};

    #[test]
    fn test_distinct_names() -> Result<(), Error> {
        let test_identity_count: usize = std::env::var_os("IDENTITY_COUNT")
            .map(|s| s.into_string().unwrap().parse().unwrap())
            .unwrap_or(16);

        let brazilian = Population {
            domain: "br",
            secret: b"0123456789abcdef0123456789abcdef",
            ingredients: &PERFUME_INGREDIENTS,
        };
        let mut store = RemoteStore {
            bridge: MockBridge::default(),
        };

        let start = Instant::now();
        let identities: Vec<Identity> = (0..test_identity_count)
            .map(|_| {
                let ident = random_hex_string::<12>();
                brazilian.identity(ident.as_str(), &mut store).unwrap()
            })
            .collect();
        let stop = Instant::now();
        println!(
            "generated {} identities in {:?}",
            identities.len(),
            stop - start
        );

        let mut pairs = vec![];
        for i in 0..identities.len() {
            for j in 0..identities.len() {
                if i < j {
                    pairs.push((i, j));
                }
            }
        }
        for (i1, i2) in &pairs {
            assert_ne!(identities[*i1].friendly_name, identities[*i2].friendly_name);
        }
        println!("compared {} names for uniqueness", pairs.len());

        let mut diff_scores: Vec<(u8, (usize, usize))> = vec![];
        for (i1, i2) in &pairs {
            let name1: Vec<&str> = identities[*i1].friendly_name.split('-').collect();
            let name2: Vec<&str> = identities[*i2].friendly_name.split('-').collect();
            match (name1.as_slice(), name2.as_slice()) {
                ([prefix1, color1, animal1], [prefix2, color2, animal2]) => {
                    let score: [u8; 3] = [
                        if prefix1 == prefix2 { 0 } else { 1 },
                        if color1 == color2 { 0 } else { 1 },
                        if animal1 == animal2 { 0 } else { 1 },
                    ];
                    diff_scores.push((score.iter().sum(), (*i1, *i2)));
                }
                _ => unimplemented!(),
            }
        }
        diff_scores.sort_by(|a, b| a.0.cmp(&b.0));

        let lowest_diff_score = diff_scores.first().unwrap().0;
        let least_diff_pairs: Vec<(u8, (usize, usize))> = diff_scores
            .into_iter()
            .take_while(|s| s.0 == lowest_diff_score)
            .collect();

        println!(
            "{} differences with {} common components",
            least_diff_pairs.len(),
            3 - lowest_diff_score
        );

        let mut least_diff_names = least_diff_pairs
            .iter()
            .map(|p| {
                let name1 = &identities[p.1.0].friendly_name;
                let name2 = &identities[p.1.1].friendly_name;
                (name1, name2)
            })
            .collect::<Vec<_>>();
        least_diff_names.sort();
        println!("most similar names: {least_diff_names:#?}",);

        Ok(())
    }
}