wikibase 0.5.0

extern crate rand;

use self::rand::prelude::SliceRandom;
use self::rand::thread_rng;
use crate::entity_diff::EntityDiff;
use crate::*;
use mediawiki::reqwest;
use std::collections::HashMap;
use std::sync::mpsc;
use std::thread;

const LOAD_ENTITIES_MAX_THREADS: usize = 10;

/// A container of `Entity` values.
/// This can load and cache entities individually, or in large batches.
/// It will load entities only once, and load groups of 50/500 entities in parallel.
#[derive(Debug, Default, Clone)]
pub struct EntityContainer {
    entities: HashMap<String, Entity>,
}

impl EntityContainer {
    /// Generates a new, empty `EntityContainer`
    pub fn new() -> EntityContainer {
        EntityContainer {
            entities: HashMap::<String, Entity>::new(),
        }
    }

    /// Loads (new) entities from the MediaWiki API, the wrapper
    pub fn load_entities(
        &mut self,
        api: &mediawiki::api::Api,
        entity_ids: &Vec<String>,
    ) -> Result<(), Box<::std::error::Error>> {
        let chunk_size = match api.user().is_bot() {
            true => 200, // Could be 500 but timeouts happen
            false => 50,
        };
        self.load_entities_internal(api, entity_ids, chunk_size)
    }

    /// Removed already existing entity IDs, removes duplicates, and shuffles the remaining ones
    pub fn unique_shuffle_entity_ids(&self, entity_ids: &Vec<String>) -> Vec<String> {
        let mut to_load = entity_ids
            .iter()
            .filter(|entity_id| !entity_id.is_empty())
            .filter(|entity_id| !self.entities.contains_key(*entity_id))
            .map(|entity_id| entity_id.to_owned())
            .collect::<Vec<String>>();

        // De-duplicate
        to_load.sort();
        to_load.dedup();

        // Shuffle entities to generate random chunks
        to_load.shuffle(&mut thread_rng());
        to_load
    }

    /// Loads (new) entities from the MediaWiki API, the actual code
    fn load_entities_internal(
        &mut self,
        api: &mediawiki::api::Api,
        entity_ids: &Vec<String>,
        chunk_size: usize,
    ) -> Result<(), Box<::std::error::Error>> {
        // Shortcut, as last resort for small chunk sizes
        if chunk_size <= 1 || entity_ids.len() == 1 {
            return self.load_entities_via_special_entity_data(api, entity_ids);
        }

        // Get list of entity IDs to actually load
        let to_load = self.unique_shuffle_entity_ids(entity_ids);

        // Something to do?
        if to_load.is_empty() {
            return Ok(());
        }

        // Create thread list
        let (tx, rx) = mpsc::channel();
        let mut thread_list = vec![];
        let mut chunks: u64 = 0;
        for chunk in to_load.chunks(chunk_size) {
            chunks = chunks + 1;
            let ids = chunk.join("|");
            let params: HashMap<_, _> = vec![
                ("action", "wbgetentities"),
                ("ids", &ids),
                ("format", "json"),
            ]
            .into_iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect();
            let req = api.get_api_request_builder(&params, "GET")?;

            let tx = mpsc::Sender::clone(&tx);
            thread_list.push((req, tx, chunk.to_owned()));
        }

        // Start initial, limited number of threads
        for _num in 0..LOAD_ENTITIES_MAX_THREADS {
            match thread_list.pop() {
                Some(x) => {
                    thread::spawn(move || {
                        x.1.send((x.0.send(), x.2))
                            .expect("Sending of result failed");
                    });
                }
                None => break,
            }
        }

        // Process each thread finishing, start a new one if any left
        let mut again: Vec<String> = vec![];
        for _ in 0..chunks {
            let thread_response = rx.recv()?;

            // Start next thread
            match thread_list.pop() {
                Some(x) => {
                    thread::spawn(move || {
                        x.1.send((x.0.send(), x.2))
                            .expect("Sending of result failed");
                    });
                }
                None => {}
            }

            let mut response = match thread_response.0 {
                Ok(response) => response,
                Err(e) => {
                    eprintln!("EntityContainer::load_entities: A chunk could not be loaded (retrying): {}",&e);
                    again.append(&mut thread_response.1.iter().cloned().collect());
                    continue;
                }
            };
            let j: serde_json::Value = match response.json() {
                Ok(j) => j,
                Err(e) => {
                    eprintln!(
                        "EntityContainer::load_entities: Could not parse JSON (retrying): {}",
                        &e
                    );
                    again.append(&mut thread_response.1.iter().cloned().collect());
                    continue;
                }
            };
            let entities = match j["entities"].as_object() {
                Some(e) => e,
                None => {
                    again.append(&mut thread_response.1.iter().cloned().collect());
                    continue;
                }
            };
            for (entity_id, entity_json) in entities {
                match self.set_entity_from_json(entity_json) {
                    Ok(_) => {}
                    Err(e) => {
                        eprintln!("Can not parse item {}: {}", &entity_id, &e);
                    }
                }
            }
        }

        // Not all chunks were loaded, retry missing ones
        // Give up if it's only one item per chunk
        if !again.is_empty() && chunk_size > 1 {
            return self.load_entities_internal(api, &again, chunk_size / 2);
        }

        Ok(())
    }

    /// Loads (new) entities from Special:EntityData
    pub fn load_entities_via_special_entity_data(
        &mut self,
        api: &mediawiki::api::Api,
        entity_ids: &Vec<String>,
    ) -> Result<(), Box<::std::error::Error>> {
        // Get list of entity IDs to actually load
        let to_load = self.unique_shuffle_entity_ids(entity_ids);

        // Something to do?
        if to_load.is_empty() {
            return Ok(());
        }

        let special_entity_data_url =
            api.get_site_info_string("general", "wikibase-conceptbaseuri")?;

        // Create thread list
        let (tx, rx) = mpsc::channel();
        let mut thread_list = vec![];
        for entity_id in &to_load {
            let url = special_entity_data_url.to_owned() + &entity_id + ".json";
            let tx = mpsc::Sender::clone(&tx);
            thread_list.push((url, tx, entity_id.to_owned()));
        }

        // Start initial, limited number of threads
        for _num in 0..LOAD_ENTITIES_MAX_THREADS {
            match thread_list.pop() {
                Some(x) => {
                    thread::spawn(move || {
                        x.1.send((reqwest::get(x.0.as_str()), x.2)).expect(
                            "load_entities_via_special_entity_data: Sending of result failed",
                        );
                    });
                }
                None => break,
            }
        }

        // Process each thread finishing, start a new one if any left
        let mut again: Vec<String> = vec![];
        for _ in 0..to_load.len() {
            let thread_response = rx.recv()?;

            // Start next thread
            match thread_list.pop() {
                Some(x) => {
                    thread::spawn(move || {
                        x.1.send((reqwest::get(x.0.as_str()), x.2)).expect(
                            "load_entities_via_special_entity_data: Sending of result failed",
                        );
                    });
                }
                None => {}
            }

            let mut response = match thread_response.0 {
                Ok(response) => response,
                Err(e) => {
                    eprintln!("EntityContainer::load_entities_via_special_entity_data: A chunk could not be loaded (retrying): {}",&e);
                    again.push(thread_response.1.to_string());
                    continue;
                }
            };
            let j: serde_json::Value = match response.json() {
                Ok(j) => j,
                Err(e) => {
                    eprintln!(
                        "EntityContainer::load_entities_via_special_entity_data: Could not parse JSON (retrying): {}",
                        &e
                    );
                    again.push(thread_response.1.to_string());
                    continue;
                }
            };
            let entities = match j["entities"].as_object() {
                Some(e) => e,
                None => {
                    again.push(thread_response.1.to_string());
                    continue;
                }
            };
            for (entity_id, entity_json) in entities {
                match self.set_entity_from_json(entity_json) {
                    Ok(_) => {}
                    Err(e) => {
                        eprintln!("Can not parse item {}: {}", &entity_id, &e);
                    }
                }
            }
        }

        if !again.is_empty() {
            eprintln!("EntityContainer::load_entities_via_special_entity_data: Entities could not be loaded: {:?}",&again);
        }

        Ok(())
    }

    /// Adds an `entity` to the cache, based on its JSON representation
    pub fn set_entity_from_json(
        &mut self,
        entity_json: &serde_json::Value,
    ) -> Result<(), Box<::std::error::Error>> {
        let entity_id = match &entity_json["id"] {
            serde_json::Value::String(s) => s.to_string(),
            _ => {
                return Err(From::from(format!(
                    "Entity has no 'id' (string) field:{}",
                    &entity_json
                )));
            }
        };
        let entity = from_json::entity_from_json(&entity_json)?;

        // Don't cache missing items
        if entity.missing() {
            self.remove_entity(entity_id);
        } else {
            self.entities.insert(entity_id.to_string(), entity);
        }
        Ok(())
    }

    /// Loads a single entity. Returns the `Entity`, or an error
    pub fn load_entity<S: Into<String>>(
        &mut self,
        api: &mediawiki::api::Api,
        entity: S,
    ) -> Result<&Entity, Box<::std::error::Error>> {
        let entity: String = entity.into();
        self.load_entities(api, &vec![entity.clone()])?;
        match self.get_entity(entity.as_str()) {
            Some(e) => Ok(e),
            None => Err(From::from(format!("No such entity '{}'", &entity))),
        }
    }

    /// Returns `Some(entity)` with that ID from the cache, or `None`.
    /// This will _not_ load entities via the API! Use `load_entity` for that
    pub fn get_entity<S: Into<String>>(&self, entity_id: S) -> Option<&Entity> {
        let entity_id: String = entity_id.into();
        self.entities.get(&entity_id)
    }

    /// Checks if an entity is in the cache.
    /// Returns true or false.
    pub fn has_entity<S: Into<String>>(&self, entity_id: S) -> bool {
        let entity_id: String = entity_id.into();
        match self.entities.get(&entity_id) {
            Some(_) => true,
            None => false,
        }
    }

    /// Removes the entity with the given key from the cache, and returns `Some(entity)` or `None`
    pub fn remove_entity<S: Into<String>>(&mut self, entity_id: S) -> Option<Entity> {
        let entity_id: String = entity_id.into();
        self.entities.remove(&entity_id)
    }

    /// Removes the entities with the given keys from the cache
    pub fn remove_entities(&mut self, entity_ids: &Vec<String>) {
        for entity_id in entity_ids {
            self.remove_entity(entity_id.to_string());
        }
    }

    /// Removes the entities with the given keys from the cache, then reloads them from the API
    pub fn reload_entities(
        &mut self,
        api: &mediawiki::api::Api,
        entity_ids: &Vec<String>,
    ) -> Result<(), Box<::std::error::Error>> {
        self.remove_entities(entity_ids);
        self.load_entities(api, entity_ids)?;
        Ok(())
    }

    /// Returns the number of cached entities
    pub fn len(&self) -> usize {
        self.entities.len()
    }

    /// Clears the cache
    pub fn clear(&mut self) {
        self.entities.clear();
    }

    /// Applies a diff, updates the cache if possible, and returns the entity ID
    pub fn apply_diff(
        &mut self,
        api: &mut mediawiki::api::Api,
        diff: &EntityDiff,
    ) -> Option<String> {
        if diff.is_empty() {
            return None; // Nothing done
        }
        match diff.apply_diff(api, &diff) {
            Ok(json) => match EntityDiff::get_entity_id(&json) {
                Some(q) => match self.set_entity_from_json(&json) {
                    Ok(_) => Some(q),
                    _ => None,
                },
                None => None,
            },
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::mediawiki::api::Api;
    use super::EntityContainer;
    use entity::*;

    /*
    // TO TEST:
    remove_entities
    reload_entities
    clear
    apply_diff
    */

    #[test]
    fn test_groups() {
        let api = Api::new("https://www.wikidata.org/w/api.php").unwrap();
        let mut ec = EntityContainer::new();

        let mut items: Vec<String> = vec![];
        for num in 1..151 {
            items.push(format!("Q{}", num));
        }

        ec.load_entities(&api, &items).unwrap();
        //println!("{} items loaded", ec.len());
        assert_eq!(ec.len(), 136);
    }

    #[test]
    fn test_remove_entity() {
        let api = Api::new("https://www.wikidata.org/w/api.php").unwrap();
        let mut ec = EntityContainer::new();
        ec.load_entities(
            &api,
            &vec!["Q42".to_string(), "Q12345".to_string(), "Q50".to_string()],
        )
        .unwrap();
        assert!(ec.has_entity("Q42"));
        assert!(ec.has_entity("Q12345"));
        ec.remove_entity("Q12345");
        assert!(ec.has_entity("Q42"));
        assert!(!ec.has_entity("Q12345"));
        ec.remove_entity("Q42");
        assert!(!ec.has_entity("Q42"));
        assert!(!ec.has_entity("Q12345"));
        assert_eq!(ec.len(), 0);
    }

    #[test]
    fn test_set_entity_from_json() {
        let j = json!({"id":"Q50","labels":{},"aliases":{},"descriptions":{},"sitelinks":{},"claims":{},"type":"item"});
        let mut ec = EntityContainer::new();
        ec.set_entity_from_json(&j).unwrap();
        assert!(ec.has_entity("Q50"));
        let entity = ec.get_entity("Q50").unwrap().to_owned();
        assert_eq!(entity.id(), "Q50");

        // Try add missing, should not add
        let j = json!({"id":"Q51","missing":""});
        ec.set_entity_from_json(&j).unwrap();
        assert!(!ec.has_entity("Q51"));
    }

    #[test]
    fn test_unique_shuffle_entity_ids() {
        let entity_ids = vec!["Q42".to_string(), "Q12345".to_string(), "Q42".to_string()];
        let ec = EntityContainer::new();
        let mut new_entity_ids = ec.unique_shuffle_entity_ids(&entity_ids);
        new_entity_ids.sort();
        assert_eq!(
            new_entity_ids,
            vec!["Q12345".to_string(), "Q42".to_string()]
        );
    }

    #[test]
    fn test_load_entity() {
        let api = Api::new("https://www.wikidata.org/w/api.php").unwrap();
        let mut ec = EntityContainer::new();
        let entity = ec.load_entity(&api, "Q12345").unwrap().to_owned();
        assert_eq!(entity.id(), "Q12345");
        assert!(ec.has_entity("Q12345"));
        let entity2 = ec.get_entity("Q12345").unwrap();
        assert_eq!(entity2.id(), "Q12345");
    }

    #[test]
    fn test_load_entities() {
        let api = Api::new("https://www.wikidata.org/w/api.php").unwrap();
        let mut ec = EntityContainer::new();
        assert_eq!(ec.len(), 0);
        ec.load_entities(
            &api,
            &vec!["Q42".to_string(), "Q12345".to_string(), "Q50".to_string()],
        )
        .unwrap();
        assert!(ec.has_entity("Q42"));
        assert!(ec.has_entity("Q12345"));
        assert!(!ec.has_entity("Q50")); // Deleted item
    }

    #[test]
    fn test_load_entities_via_special_entity_data() {
        let api = Api::new("https://www.wikidata.org/w/api.php").unwrap();
        let mut ec = EntityContainer::new();
        ec.load_entities_via_special_entity_data(
            &api,
            &vec!["Q42".to_string(), "Q12345".to_string(), "Q50".to_string()],
        )
        .unwrap();
        assert!(ec.has_entity("Q42"));
        assert!(ec.has_entity("Q12345"));
        assert!(!ec.has_entity("Q50")); // Deleted item
    }
}