mwapi 0.5.0-alpha.2

/*
Copyright (C) 2021-2022 Kunal Mehta <legoktm@debian.org>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use crate::client::InnerClient;
use crate::{Client, Error, Method, Params, Result};
use indexmap::IndexMap;
use mwapi_responses::normalize;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
use tokio::sync::oneshot;
use tokio::sync::Mutex;
use tokio::time;
use tracing::debug;

/// Set of query params that uniquely identify a request
type Request = BTreeMap<String, String>;
type Sender = oneshot::Sender<Result<Value>>;

/// Check whether the request is combinable
pub(crate) fn is_combinable(params: &Params) -> bool {
    // Must be action=query
    if params.get("action") != Some(&"query".to_string()) {
        return false;
    }

    // Must have a single-valued titles parameter
    // TODO: support multi-valued titles
    if let Some(titles) = params.get("titles") {
        if titles.contains('|') {
            return false;
        }
    } else {
        return false;
    }

    // For now we're only merging prop=
    if params.contains_key("list")
        || params.contains_key("meta")
        || params.contains_key("generator")
        // Also we can't support &redirects=1 yet
        || params.contains_key("redirects")
    {
        return false;
    }

    true
}

#[derive(Default, Debug)]
pub(crate) struct RequestQueue {
    queue: Mutex<IndexMap<Request, HashMap<String, Sender>>>,
}

impl RequestQueue {
    pub(crate) fn init(&self, client: &Client) {
        let client = client.inner.clone();
        tokio::spawn(async move { background_thread(client).await });
    }

    pub(crate) async fn push(&self, params: Params, sender: Sender) {
        let mut params = params.map;
        // unwrap: is_combinable() checks that a titles parameter is present
        let title = params.remove("titles").unwrap();
        let req: Request = params.into_iter().collect();
        self.queue
            .lock()
            .await
            .entry(req)
            .or_default()
            .insert(title, sender);
    }

    pub(crate) async fn pop(
        &self,
    ) -> Option<(Request, HashMap<String, Sender>)> {
        self.queue.lock().await.pop()
    }
}

#[derive(Deserialize)]
struct Response {
    query: QueryResponse,
}

#[derive(Deserialize)]
struct QueryResponse {
    #[serde(default)]
    normalized: Vec<normalize::Normalized>,
    #[serde(default)]
    redirects: Vec<normalize::Redirect>,
    #[serde(default)]
    pages: Vec<PageResponse>,
}

#[derive(Deserialize, Serialize)]
struct PageResponse {
    title: String,
    #[serde(flatten)]
    extra: HashMap<String, Value>,
}

/// The title_map is user input -> normalized. Except we want to go from
/// the normalized title to user input to match with the correct sender.
/// So get the title map out of the response and flip it.
///
/// This is literally mwapi_responses::ApiResponse::title_map() but backwards
fn flipped_title_map(resp: &QueryResponse) -> HashMap<String, String> {
    let redirects: HashMap<String, String> = resp
        .redirects
        .iter()
        .map(|redir| (redir.to.to_string(), redir.from.to_string()))
        .collect();
    let mut normalized: HashMap<String, String> = resp
        .normalized
        .iter()
        .map(|norm| (norm.to.to_string(), norm.from.to_string()))
        .collect();
    let mut titles = HashMap::new();
    for (to, from) in redirects {
        // See if the redirect is also a normalized title
        let real_from = normalized.remove(&from).unwrap_or(from);
        titles.insert(to, real_from);
    }
    // Merge in the remaining normalized titles
    titles.extend(normalized);

    titles
}

/// Run a background thread that checks for new pending requests,
/// processes them, in batches if needed, and sends the result back
/// to various receivers.
///
/// This thread very carefully only holds onto a weak reference of InnerClient
/// to avoid circular dependencies when the client is actually dropped.
pub(crate) async fn background_thread(client: Arc<InnerClient>) {
    let weak_client = Arc::downgrade(&client);
    // We specifically drop the strong reference to client here to break
    // the dependency loop, we only want to hold onto a weak reference.
    drop(client);
    let mut interval = time::interval(time::Duration::from_micros(50_000));
    loop {
        let client = match weak_client.upgrade() {
            Some(client) => client,
            None => {
                // All references to InnerClient have been dropped, so we
                // can abort now.
                return;
            }
        };
        let next = client.get_queue.pop().await;
        let (params, mut senders) = match next {
            Some(val) => val,
            None => {
                // Small sleep to avoid holding the lock in a tight loop
                interval.tick().await;
                continue;
            }
        };
        let titles: Vec<_> = senders.keys().map(|t| t.to_string()).collect();
        debug!("queue: {:?}", &titles);
        // limit for titles is 50/500. Let's do 25 to be on the safe side
        for chunk in titles.chunks(25) {
            let mut params2 = params.clone();
            params2.insert("titles".to_string(), chunk.join("|"));
            let result: Result<Response> =
                client.request(Method::Get, params2).await;
            match result {
                Ok(response) => {
                    let title_map = flipped_title_map(&response.query);
                    for page in response.query.pages {
                        let title = page.title.to_string();
                        // Use the title map to get the input we provided, not
                        // the normalized form
                        let unnormalized = match title_map.get(&title) {
                            Some(title) => title.to_string(),
                            None => title.to_string(),
                        };
                        if let Some(sender) = senders.remove(&unnormalized) {
                            // TODO: normalized/redirects
                            let split = serde_json::json!({
                                "batchcomplete": true,
                                "query": {
                                    "pages": [page]
                                }
                            });
                            if sender.send(Ok(split)).is_err() {
                                debug!("{} dropped receiver", &title);
                            }
                        }
                    }
                }
                Err(err) => {
                    for title in chunk {
                        if let Some(sender) = senders.remove(title) {
                            if sender.send(Err(err.clone())).is_err() {
                                debug!("{} dropped receiver", &title);
                            }
                        }
                    }
                }
            }
        }
        // Drop the strong reference, we don't need it any more
        drop(client);

        // If there are any senders left, error out
        for (title, sender) in senders {
            if sender
                .send(Err(Error::Unknown(
                    "Couldnt find page response".to_string(),
                )))
                .is_err()
            {
                debug!(
                    "{} couldnt be found in response, and sender failed",
                    &title
                );
            }
        }
    }
}