use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, instrument, warn};
use super::Page;
use super::locator::{Locator, Selector};
use crate::error::LocatorError;
pub type LocatorHandlerFn =
Arc<dyn Fn() -> Pin<Box<dyn Future<Output = Result<(), LocatorError>> + Send>> + Send + Sync>;
#[derive(Debug, Clone, Default)]
pub struct LocatorHandlerOptions {
pub no_wait_after: bool,
pub times: Option<u32>,
}
#[derive(Clone)]
struct RegisteredHandler {
id: u64,
selector: Selector,
handler: LocatorHandlerFn,
options: LocatorHandlerOptions,
run_count: u32,
}
impl std::fmt::Debug for RegisteredHandler {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RegisteredHandler")
.field("id", &self.id)
.field("selector", &self.selector)
.field("options", &self.options)
.field("run_count", &self.run_count)
.finish()
}
}
#[derive(Debug)]
pub struct LocatorHandlerManager {
handlers: RwLock<Vec<RegisteredHandler>>,
next_id: std::sync::atomic::AtomicU64,
}
impl LocatorHandlerManager {
pub fn new() -> Self {
Self {
handlers: RwLock::new(Vec::new()),
next_id: std::sync::atomic::AtomicU64::new(1),
}
}
#[instrument(level = "debug", skip(self, handler), fields(selector = ?selector))]
pub async fn add_handler<F, Fut>(
&self,
selector: Selector,
handler: F,
options: LocatorHandlerOptions,
) -> u64
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<(), LocatorError>> + Send + 'static,
{
let id = self
.next_id
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let handler_fn: LocatorHandlerFn = Arc::new(move || Box::pin(handler()));
let registered = RegisteredHandler {
id,
selector,
handler: handler_fn,
options,
run_count: 0,
};
let mut handlers = self.handlers.write().await;
handlers.push(registered);
debug!(handler_id = id, "Locator handler registered");
id
}
#[instrument(level = "debug", skip(self))]
pub async fn remove_handler_by_id(&self, id: u64) {
let mut handlers = self.handlers.write().await;
let initial_len = handlers.len();
handlers.retain(|h| h.id != id);
if handlers.len() < initial_len {
debug!(handler_id = id, "Locator handler removed");
} else {
debug!(handler_id = id, "No matching locator handler found");
}
}
#[instrument(level = "debug", skip(self, page))]
pub async fn try_handle_blocking(&self, page: &Page) -> bool {
let handlers = self.handlers.read().await;
for handler in handlers.iter() {
let locator = Locator::new(page, handler.selector.clone());
if let Ok(is_visible) = locator.is_visible().await {
if is_visible {
let handler_id = handler.id;
debug!(
handler_id = handler_id,
"Handler selector matched, running handler"
);
let handler_fn = handler.handler.clone();
drop(handlers);
if let Err(e) = handler_fn().await {
warn!(handler_id = handler_id, "Locator handler failed: {}", e);
} else {
let mut handlers = self.handlers.write().await;
if let Some(handler) = handlers.iter_mut().find(|h| h.id == handler_id) {
handler.run_count += 1;
if let Some(times) = handler.options.times {
if handler.run_count >= times {
debug!(
handler_id = handler_id,
"Handler reached times limit, removing"
);
handlers.retain(|h| h.id != handler_id);
}
}
}
return true;
}
return false;
}
}
}
false
}
}
impl Default for LocatorHandlerManager {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Copy)]
pub struct LocatorHandlerHandle {
id: u64,
}
impl LocatorHandlerHandle {
pub(crate) fn new(id: u64) -> Self {
Self { id }
}
pub fn id(&self) -> u64 {
self.id
}
}
impl super::Page {
pub async fn add_locator_handler<F, Fut>(
&self,
locator: impl Into<super::Locator<'_>>,
handler: F,
) -> LocatorHandlerHandle
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = Result<(), crate::error::LocatorError>> + Send + 'static,
{
let loc = locator.into();
let id = self
.locator_handler_manager
.add_handler(
loc.selector().clone(),
handler,
LocatorHandlerOptions::default(),
)
.await;
LocatorHandlerHandle::new(id)
}
pub async fn add_locator_handler_with_options<F, Fut>(
&self,
locator: impl Into<super::Locator<'_>>,
handler: F,
options: LocatorHandlerOptions,
) -> LocatorHandlerHandle
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = Result<(), crate::error::LocatorError>> + Send + 'static,
{
let loc = locator.into();
let id = self
.locator_handler_manager
.add_handler(loc.selector().clone(), handler, options)
.await;
LocatorHandlerHandle::new(id)
}
pub async fn remove_locator_handler(&self, handle: LocatorHandlerHandle) {
self.locator_handler_manager
.remove_handler_by_id(handle.id())
.await;
}
pub(crate) fn locator_handler_manager(&self) -> &LocatorHandlerManager {
&self.locator_handler_manager
}
}