wikibase 0.7.6

use crate::entity_diff::EntityDiff;
use crate::*;
use futures::future::join_all;
use futures::StreamExt;
use mediawiki::reqwest;
use rand::prelude::*;
use rayon::prelude::*;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::{thread, time};

//const LOAD_ENTITIES_MAX_THREADS: usize = 10;
const DELAY_BEFORE_TRYING_AGAIN_MS: u64 = 500;

/// 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: Arc<RwLock<HashMap<String, Entity>>>,
    special_entity_data_allowed: bool,
    special_entity_data_revision_attempts: usize,
    max_concurrent: Option<usize>,
}

impl EntityContainer {
    /// Generates a new, empty `EntityContainer`
    pub fn new() -> EntityContainer {
        EntityContainer {
            entities: Arc::new(RwLock::new(HashMap::<String, Entity>::new())),
            special_entity_data_allowed: false, // The cache lag is horrible
            special_entity_data_revision_attempts: 5,
            max_concurrent: None,
        }
    }

    pub fn set_max_concurrent(&mut self, max_concurrent: usize) {
        self.max_concurrent = Some(max_concurrent);
    }

    pub fn unset_max_concurrent(&mut self) {
        self.max_concurrent = None;
    }

    pub fn get_max_concurrent(&self) -> &Option<usize> {
        &self.max_concurrent
    }

    /// Allows fallback to [[Special:EntityData]]. That is faster but can lag behind newly created items
    pub fn allow_special_entity_data(&mut self, allow: bool) {
        self.special_entity_data_allowed = allow;
    }

    /// Sets the number of attempts for loading a specific revision via [[Special:EntityData]]
    pub fn set_special_entity_data_revision_attempts(&mut self, num: usize) {
        self.special_entity_data_revision_attempts = num;
    }

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

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

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

        // Shuffle entities to generate random chunks
        to_load.shuffle(&mut rand::rng());
        Ok(to_load)
    }

    /*
    fn get_thread_pool(&self) -> Result<ThreadPool, rayon::ThreadPoolBuildError> {
        let mut attempts_left: u32 = 5;
        loop {
            match ThreadPoolBuilder::new()
                //.num_threads(LOAD_ENTITIES_MAX_THREADS)
                .build()
            {
                Ok(ret) => return Ok(ret),
                Err(e) => {
                    if attempts_left == 0 {
                        return Err(e);
                    }
                    let delay = time::Duration::from_millis(50);
                    thread::sleep(delay);
                    attempts_left -= 1;
                }
            }
        }
    }
    */

    /// Loads (new) entities from the MediaWiki API, the actual code
    async fn load_entities_internal(
        &self,
        api: &mediawiki::api::Api,
        entity_ids: &Vec<String>,
        chunk_size: usize,
    ) -> Result<(), WikibaseError> {
        let mut chunk_size = chunk_size;

        loop {
            // Shortcut, as last resort for small chunk sizes
            if self.special_entity_data_allowed && (chunk_size <= 1 || entity_ids.len() == 1) {
                return self
                    .load_entities_via_special_entity_data(api, entity_ids)
                    .await;
            }

            if chunk_size < 1 {
                return Err(WikibaseError::ChunkSizeZero);
            }

            // 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 params_list: Vec<_> = to_load
                .chunks(chunk_size)
                .map(|chunk| chunk.join("|"))
                .map(|ids| {
                    api.params_into(&[
                        ("action", "wbgetentities"),
                        ("ids", &ids),
                        ("format", "json"),
                    ])
                })
                .collect();

            let futures: Vec<_> = params_list
                .iter()
                .map(|params| api.get_query_api_json(params))
                .collect();

            let results = match self.max_concurrent {
                Some(max_concurrent) => {
                    let stream = futures::stream::iter(futures).buffer_unordered(max_concurrent);
                    stream.collect::<Vec<_>>().await
                }
                None => join_all(futures).await,
            };

            let again: Vec<String> = to_load
                .iter()
                .zip(results)
                .filter_map(|(ids, result)| {
                    let j = match result {
                        Ok(j) => j,
                        _ => return Some(ids),
                    };

                    let entities = match j["entities"].as_object() {
                        Some(e) => e,
                        None => return Some(ids),
                    };

                    let bad_entities: Vec<String> = entities
                        .iter()
                        .filter_map(|(entity_id, entity_json)| {
                            match self.set_entity_from_json(entity_json) {
                                Ok(_) => None,
                                Err(_) => Some(entity_id.to_owned()),
                            }
                        })
                        .collect();

                    if !bad_entities.is_empty() {
                        eprintln!("Can not parse item(s) {:?}", &bad_entities);
                    }
                    None
                })
                .cloned()
                .collect();

            let again: Vec<String> = again
                .iter()
                .flat_map(|x| x.split('|').map(|s| s.to_string()).collect::<Vec<String>>())
                .collect();

            if !again.is_empty() && chunk_size > 1 {
                // Some items might be probably too new; wait a bit
                let delay = time::Duration::from_millis(DELAY_BEFORE_TRYING_AGAIN_MS);
                thread::sleep(delay);
                //return self.load_entities_internal(api, &again, chunk_size / 2).await ;
                chunk_size /= 2;
            } else {
                return Ok(());
            }
        }
    }

    /// Loads (new) entities from Special:EntityData
    pub async fn load_entities_via_special_entity_data(
        &self,
        api: &mediawiki::api::Api,
        entity_ids: &Vec<String>,
    ) -> Result<(), WikibaseError> {
        // 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")?;

        let mut again: Vec<String> = vec![];
        for entity_id in to_load {
            let url = format!("{}{}.json", &special_entity_data_url, entity_id);
            let response = match reqwest::get(url.as_str()).await {
                Ok(r) => r,
                _ => continue,
            };
            let j: serde_json::Value = match response.json().await {
                Ok(e) => e,
                _ => {
                    again.push(entity_id);
                    continue;
                }
            };

            let entities = match j["entities"].as_object() {
                Some(e) => e,
                None => {
                    again.push(entity_id);
                    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() {
            // Some items might be probably too new; wait a bit
            let delay = time::Duration::from_millis(DELAY_BEFORE_TRYING_AGAIN_MS);
            thread::sleep(delay);
            eprintln!("EntityContainer::load_entities_via_special_entity_data: Entities could not be loaded: {:?}",&again);
        }

        Ok(())
    }

    /// Loads a specific revision of an entity from Special:EntityData
    pub async fn load_entity_revision_via_special_entity_data(
        &self,
        api: &mediawiki::api::Api,
        entity_id: String,
        revision_id: usize,
    ) -> Result<(), WikibaseError> {
        let special_entity_data_url =
            api.get_site_info_string("general", "wikibase-conceptbaseuri")?;

        let mut attempts_left = self.special_entity_data_revision_attempts;

        let url = format!(
            "{}{}.json?revision={}",
            &special_entity_data_url, &entity_id, &revision_id,
        );

        loop {
            if attempts_left == 0 {
                return Err(WikibaseError::Request(format!("EntityContainer::load_entity_revision_via_special_entity_data: Entity revision could not be loaded after multiple attempts: {}/{}",&entity_id,&revision_id)));
            }
            attempts_left -= 1;
            let response = reqwest::get(url.as_str()).await;
            let response = match response {
                Ok(response) => response,
                Err(e) => {
                    return Err(WikibaseError::Request(format!("EntityContainer::load_entity_revision_via_special_entity_data: Entity revision could not be loaded: {}/{}: {}",&entity_id,&revision_id,&e)));
                }
            };
            let entity_json: serde_json::Value = match response.json().await {
                Ok(j) => j,
                Err(_e) => {
                    // Revision is probably too new; wait a bit
                    let delay = time::Duration::from_millis(DELAY_BEFORE_TRYING_AGAIN_MS);
                    thread::sleep(delay);
                    continue;
                }
            };
            let entity_json = &entity_json["entities"][&entity_id];

            match self.set_entity_from_json(entity_json) {
                Ok(_) => return Ok(()),
                Err(e) => {
                    return Err(WikibaseError::Serialization(format!("EntityContainer::load_entity_revision_via_special_entity_data: Could not parse entity JSON: {}/{}: {}",&entity_id,&revision_id,&e)));
                }
            }
        }
    }

    /// Adds an `entity` to the cache, based on its JSON representation
    pub fn set_entity_from_json(
        &self,
        entity_json: &serde_json::Value,
    ) -> Result<(), WikibaseError> {
        let entity_id = match &entity_json["id"] {
            serde_json::Value::String(s) => s.to_string(),
            _ => {
                return Err(WikibaseError::Serialization(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 if let Ok(mut entities) = self.entities.write() {
            entities.insert(entity_id.to_string(), entity);
        }

        Ok(())
    }

    /// Adds a claim (from wbcreateclaim) to an entity
    pub fn add_claim_to_entity<S: Into<String>>(
        &self,
        entity: S,
        claim_json: &serde_json::Value,
    ) -> Result<(), WikibaseError> {
        let entity: String = entity.into();
        let mut entities = match self.entities.write() {
            Ok(e) => e,
            _ => {
                return Err(WikibaseError::PoisonedMutex);
            }
        };
        match entities.get_mut(&entity) {
            Some(entity) => {
                let claim = Statement::new_from_json(claim_json)?;
                entity.claims_mut().push(claim);
                Ok(())
            }
            None => Err(WikibaseError::Validation(format!(
                "Entity {entity} not in container"
            ))),
        }
    }

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

    /// Loads a specific revision of a single entity. Returns the `Entity`, or an error
    pub async fn load_entity_revision<S: Into<String>>(
        &self,
        api: &mediawiki::api::Api,
        entity: S,
        revision_id: Option<usize>,
    ) -> Result<Entity, WikibaseError> {
        match revision_id {
            Some(revision_id) => {
                let entity: String = entity.into();
                self.load_entity_revision_via_special_entity_data(
                    api,
                    entity.to_owned(),
                    revision_id,
                )
                .await?;
                match self.get_entity(entity.as_str()) {
                    Some(e) => Ok(e),
                    None => Err(WikibaseError::Validation(format!(
                        "No such entity/revision '{entity}/{revision_id}'"
                    ))),
                }
            }
            None => self.load_entity(api, entity).await,
        }
    }

    /// 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> {
        match self.entities.read() {
            Ok(e) => e.get(&entity_id.into()).map(|x| x.to_owned()),
            _ => None,
        }
    }

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

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

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

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

    /// Returns the number of cached entities
    pub fn len(&self) -> usize {
        match self.entities.read() {
            Ok(entities) => entities.len(),
            _ => 0,
        }
    }

    /// Returns true if the cache is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Clears the cache
    pub fn clear(&self) {
        if let Ok(mut entities) = self.entities.write() {
            entities.clear()
        }
    }

    /// Applies a diff, updates the cache if possible, and returns the entity ID
    pub async fn apply_diff(
        &self,
        api: &mut mediawiki::api::Api,
        diff: &EntityDiff,
    ) -> Option<String> {
        if diff.is_empty() {
            return None; // Nothing done
        }
        match diff.apply_diff(api, diff).await {
            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 crate::entity::*;
    use crate::test_helpers::*;
    use wiremock::matchers::{path, query_param};
    use wiremock::{Mock, ResponseTemplate};

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

    #[tokio::test]
    async fn test_groups() {
        let server = start_wikidata_mock().await;
        Mock::given(query_param("action", "wbgetentities"))
            .respond_with(ResponseTemplate::new(200).set_body_json(q1_to_q150_response()))
            .mount(&server)
            .await;
        let api = Api::new(&server.uri()).await.unwrap();
        let ec = EntityContainer::new();

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

        ec.load_entities(&api, &items).await.unwrap();
        // Q10, Q20, ..., Q140 (14 items) are mocked as missing; 136 valid
        assert_eq!(ec.len(), 136);
    }

    #[tokio::test]
    async fn test_remove_entity() {
        let server = start_wikidata_mock().await;
        Mock::given(query_param("action", "wbgetentities"))
            .respond_with(
                ResponseTemplate::new(200).set_body_json(q42_q12345_q50_response()),
            )
            .mount(&server)
            .await;
        let api = Api::new(&server.uri()).await.unwrap();
        let ec = EntityContainer::new();
        ec.load_entities(
            &api,
            &vec!["Q42".to_string(), "Q12345".to_string(), "Q50".to_string()],
        )
        .await
        .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);
    }

    #[tokio::test]
    async fn test_set_entity_from_json() {
        let j = json!({"id":"Q50","labels":{},"aliases":{},"descriptions":{},"sitelinks":{},"claims":{},"type":"item"});
        let 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"));
    }

    #[tokio::test]
    async 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).unwrap();
        new_entity_ids.sort();
        assert_eq!(
            new_entity_ids,
            vec!["Q12345".to_string(), "Q42".to_string()]
        );
    }

    #[tokio::test]
    async fn test_load_entity() {
        let server = start_wikidata_mock().await;
        Mock::given(query_param("action", "wbgetentities"))
            .respond_with(
                ResponseTemplate::new(200)
                    .set_body_json(wbgetentities_response(vec![q12345_entity()])),
            )
            .mount(&server)
            .await;
        let api = Api::new(&server.uri()).await.unwrap();
        let ec = EntityContainer::new();
        let entity = ec.load_entity(&api, "Q12345").await.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");
    }

    #[tokio::test]
    async fn test_load_entities() {
        let server = start_wikidata_mock().await;
        Mock::given(query_param("action", "wbgetentities"))
            .respond_with(
                ResponseTemplate::new(200).set_body_json(q42_q12345_q50_response()),
            )
            .mount(&server)
            .await;
        let api = Api::new(&server.uri()).await.unwrap();
        let ec = EntityContainer::new();
        assert_eq!(ec.len(), 0);
        ec.load_entities(
            &api,
            &vec!["Q42".to_string(), "Q12345".to_string(), "Q50".to_string()],
        )
        .await
        .unwrap();
        assert!(ec.has_entity("Q42"));
        assert!(ec.has_entity("Q12345"));
        assert!(!ec.has_entity("Q50")); // Mocked as missing
    }

    #[tokio::test]
    async fn test_load_entities_via_special_entity_data() {
        let server = start_wikidata_mock().await;
        Mock::given(path("/entity/Q42.json"))
            .respond_with(
                ResponseTemplate::new(200)
                    .set_body_json(wbgetentities_response(vec![q42_entity()])),
            )
            .mount(&server)
            .await;
        Mock::given(path("/entity/Q12345.json"))
            .respond_with(
                ResponseTemplate::new(200)
                    .set_body_json(wbgetentities_response(vec![q12345_entity()])),
            )
            .mount(&server)
            .await;
        Mock::given(path("/entity/Q50.json"))
            .respond_with(
                ResponseTemplate::new(200)
                    .set_body_json(wbgetentities_response(vec![missing_entity("Q50")])),
            )
            .mount(&server)
            .await;
        let api = Api::new(&server.uri()).await.unwrap();
        let ec = EntityContainer::new();
        ec.load_entities_via_special_entity_data(
            &api,
            &vec!["Q42".to_string(), "Q12345".to_string(), "Q50".to_string()],
        )
        .await
        .unwrap();
        assert!(ec.has_entity("Q42"));
        assert!(ec.has_entity("Q12345"));
        assert!(!ec.has_entity("Q50")); // Mocked as missing
    }

    #[tokio::test]
    async fn test_load_entity_revision_via_special_entity_data() {
        let server = start_wikidata_mock().await;
        // Specific mocks registered first (first-match-wins in wiremock)
        // Non-existing revision for Q42 — returns empty entities so set_entity_from_json fails
        Mock::given(path("/entity/Q42.json"))
            .and(query_param("revision", "12345"))
            .respond_with(
                ResponseTemplate::new(200)
                    .set_body_json(serde_json::json!({"entities": {}, "success": 1})),
            )
            .mount(&server)
            .await;
        // Non-existing revision for Q50
        Mock::given(path("/entity/Q50.json"))
            .and(query_param("revision", "12345"))
            .respond_with(
                ResponseTemplate::new(200)
                    .set_body_json(serde_json::json!({"entities": {}, "success": 1})),
            )
            .mount(&server)
            .await;
        // Generic Q42 path (matches Q42 with any or no revision not matched above)
        Mock::given(path("/entity/Q42.json"))
            .respond_with(
                ResponseTemplate::new(200)
                    .set_body_json(wbgetentities_response(vec![q42_entity()])),
            )
            .mount(&server)
            .await;
        let api = Api::new(&server.uri()).await.unwrap();

        // Existing revision
        let ec = EntityContainer::new();
        ec.load_entity_revision_via_special_entity_data(&api, "Q42".to_string(), 988117536)
            .await
            .unwrap();
        assert!(ec.has_entity("Q42"));

        // Existing item, non-existing revision
        let mut ec = EntityContainer::new();
        ec.special_entity_data_revision_attempts = 1;
        if ec
            .load_entity_revision_via_special_entity_data(&api, "Q42".to_string(), 12345)
            .await
            .is_ok()
        {
            panic!("Something's wrong")
        }
        assert!(!ec.has_entity("Q42"));

        // Non-existing item, non-existing revision
        let mut ec = EntityContainer::new();
        ec.special_entity_data_revision_attempts = 1;
        if ec
            .load_entity_revision_via_special_entity_data(&api, "Q50".to_string(), 12345)
            .await
            .is_ok()
        {
            panic!("Something's wrong")
        }
        assert!(!ec.has_entity("Q50"));
    }
}