nekosbest 0.21.6

nekos.best API wrapper
Documentation
use std::ops::Add;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::time::Instant;

#[cfg(not(feature = "blocking"))]
pub(crate) type ReqwestClient = reqwest::Client;
#[cfg(feature = "blocking")]
pub(crate) type ReqwestClient = reqwest::blocking::Client;

#[cfg(not(feature = "blocking"))]
pub(crate) type ReqBuilder = reqwest::RequestBuilder;
#[cfg(feature = "blocking")]
pub(crate) type ReqBuilder = reqwest::blocking::RequestBuilder;

#[cfg(not(feature = "blocking"))]
pub(crate) type ReqwestResponse = reqwest::Response;
#[cfg(feature = "blocking")]
pub(crate) type ReqwestResponse = reqwest::blocking::Response;

struct SearchRatelimitData {
    remaining: u32,
    resets_at: Instant,
}

#[derive(Copy, Clone, Default)]
pub enum SearchRatelimitBehavior {
    #[default]
    Sleep,
    Error,
}

#[derive(Default, Copy, Clone)]
pub struct ClientConfig {
    pub search_ratelimit_behavior: SearchRatelimitBehavior,
}

pub struct Client {
    pub(crate) client: ReqwestClient,
    client_config: ClientConfig,
    search_ratelimit_data: Arc<Mutex<Option<SearchRatelimitData>>>,
}

impl Client {
    pub fn new(client_config: ClientConfig) -> Self {
        Self::new_with_reqwest_client(ReqwestClient::new(), client_config)
    }

    pub fn new_with_reqwest_client(client: ReqwestClient, client_config: ClientConfig) -> Self {
        Self {
            client,
            client_config,
            search_ratelimit_data: Arc::new(Mutex::new(None)),
        }
    }

    pub(crate) async fn handle_search_ratelimit(&self) -> Result<(), crate::NekosBestError> {
        let lock = self.search_ratelimit_data.lock().await;
        if let Some(search_ratelimit_data) = &*lock
            && search_ratelimit_data.remaining == 0
        {
            match self.client_config.search_ratelimit_behavior {
                SearchRatelimitBehavior::Sleep => {
                    tokio::time::sleep_until(search_ratelimit_data.resets_at).await;
                }
                SearchRatelimitBehavior::Error => {
                    return Err(crate::NekosBestError::RateLimited);
                }
            }
        }

        Ok(())
    }

    pub(crate) async fn update_search_ratelimit_data(&self, headers: &reqwest::header::HeaderMap) {
        let Some(remaining) = headers.get("X-Rate-Limit-Remaining") else {
            return;
        };
        let Some(reset) = headers.get("X-Rate-Limit-Reset") else {
            return;
        };
        let Ok(remaining) = remaining.to_str() else {
            return;
        };
        let Ok(remaining) = remaining.parse::<u32>() else {
            return;
        };
        let Ok(reset) = reset.to_str() else {
            return;
        };

        let Ok(reset) = reset.parse::<chrono::DateTime<chrono::Utc>>() else {
            return;
        };

        let diff = reset - chrono::Utc::now();

        if diff < chrono::Duration::zero() {
            *self.search_ratelimit_data.lock().await = None;
            return;
        }

        *self.search_ratelimit_data.lock().await = Some(SearchRatelimitData {
            remaining,
            resets_at: Instant::from_std(std::time::Instant::now().add(diff.to_std().unwrap())),
        });
    }
}