botcore/

types.rs

use crate::error::Result;
use async_trait::async_trait;
use std::pin::Pin;
use tokio_stream::Stream;
use tokio_stream::StreamExt;

/// A stream of events emitted by a [Collector].
///
/// This type alias represents an asynchronous stream of events that can be
/// consumed by the engine. The stream is boxed and pinned to allow for
/// dynamic dispatch and async iteration.
pub type CollectorStream<'a, E> = Pin<Box<dyn Stream<Item = E> + Send + 'a>>;

/// A source of events that can be processed by the engine.
///
/// Collectors are responsible for producing events from external sources like:
/// - Blockchain events
/// - API webhooks
/// - Database changes
/// - File system events
/// - etc.
///
/// # Example
///
/// ```rust
/// use botcore::types::{Collector, CollectorStream};
/// use botcore::error::Result;
/// use async_trait::async_trait;
///
/// struct BlockEvent {
///     block_number: u64,
/// }
///
/// struct BlockCollector {
///     ws_endpoint: String,
/// }
///
/// #[async_trait]
/// impl Collector<BlockEvent> for BlockCollector {
///     async fn get_event_stream(&self) -> Result<CollectorStream<'_, BlockEvent>> {
///         // Implementation to stream block events
///         # todo!()
///     }
/// }
/// ```
#[async_trait]
pub trait Collector<E>: Send + Sync {
    /// Returns the core event stream for the collector.
    ///
    /// This method should establish any necessary connections and return a stream
    /// that will emit events of type `E`. The stream should be infinite unless
    /// an error occurs.
    ///
    /// # Errors
    ///
    /// This method should return an error if:
    /// - The connection to the event source fails
    /// - Required resources are unavailable
    /// - The stream cannot be established for any other reason
    async fn get_event_stream(&self) -> Result<CollectorStream<'_, E>>;
}

/// A component that processes events and generates actions.
///
/// Strategies contain the core business logic of the bot. They receive events
/// from collectors, process them according to some rules or algorithms, and
/// optionally generate actions to be executed.
///
/// # Type Parameters
///
/// * `E` - The type of events this strategy can process
/// * `A` - The type of actions this strategy can generate
///
/// # Example
///
/// ```rust
/// use botcore::types::Strategy;
/// use botcore::error::Result;
/// use async_trait::async_trait;
///
/// struct BlockEvent {
///     block_number: u64,
/// }
///
/// struct TradeAction {
///     amount: u64,
/// }
///
/// struct TradingStrategy {
///     min_block_interval: u64,
///     last_trade_block: u64,
/// }
///
/// #[async_trait]
/// impl Strategy<BlockEvent, TradeAction> for TradingStrategy {
///     async fn sync_state(&mut self) -> Result<()> {
///         // Sync last trade block from chain
///         # Ok(())
///     }
///
///     async fn process_event(&mut self, event: BlockEvent) -> Vec<TradeAction> {
///         if event.block_number >= self.last_trade_block + self.min_block_interval {
///             vec![TradeAction { amount: 100 }]
///         } else {
///             vec![]
///         }
///     }
/// }
/// ```
#[async_trait]
pub trait Strategy<E, A>: Send + Sync {
    /// Synchronizes the initial state of the strategy.
    ///
    /// This method is called once when the strategy is started. It should:
    /// - Load any necessary state from storage
    /// - Initialize connections to required services
    /// - Perform any other one-time setup
    ///
    /// # Errors
    ///
    /// This method should return an error if the state cannot be synchronized.
    async fn sync_state(&mut self) -> Result<()>;

    /// Processes an event and optionally generates actions.
    ///
    /// This is the core method where the strategy's business logic lives.
    /// It receives events one at a time and can generate zero or more actions
    /// in response.
    ///
    /// # Arguments
    ///
    /// * `event` - The event to process
    ///
    /// # Returns
    ///
    /// A vector of actions to be executed. An empty vector means no actions
    /// should be taken for this event.
    async fn process_event(&mut self, event: E) -> Vec<A>;
}

/// A component that executes actions generated by strategies.
///
/// Executors are responsible for carrying out the actions decided upon by
/// strategies. This might involve:
/// - Submitting transactions
/// - Making API calls
/// - Writing to databases
/// - Sending notifications
/// - etc.
///
/// # Example
///
/// ```rust
/// use botcore::types::Executor;
/// use botcore::error::Result;
/// use async_trait::async_trait;
///
/// struct TradeAction {
///     amount: u64,
/// }
///
/// struct TradeExecutor {
///     api_key: String,
/// }
///
/// #[async_trait]
/// impl Executor<TradeAction> for TradeExecutor {
///     async fn execute(&self, action: TradeAction) -> Result<()> {
///         // Execute trade via API
///         # Ok(())
///     }
/// }
/// ```
#[async_trait]
pub trait Executor<A>: Send + Sync {
    /// Executes a single action.
    ///
    /// This method should handle all the logic needed to execute the action,
    /// including any retries or error handling.
    ///
    /// # Arguments
    ///
    /// * `action` - The action to execute
    ///
    /// # Errors
    ///
    /// This method should return an error if the action cannot be executed
    /// successfully after all retries are exhausted.
    async fn execute(&self, action: A) -> Result<()>;
}

/// A wrapper around a [Collector] that maps outgoing events to a different type.
///
/// This adapter allows you to transform events from one type to another as they
/// flow through the collector. This is useful when you need to:
/// - Convert between different event representations
/// - Filter out unwanted events
/// - Enrich events with additional data
///
/// # Type Parameters
///
/// * `E` - The original event type
/// * `F` - The function type that maps events
pub struct CollectorMap<E, F> {
    collector: Box<dyn Collector<E>>,
    f: F,
}

impl<E, F> CollectorMap<E, F> {
    /// Creates a new collector map with the given collector and mapping function.
    pub fn new(collector: Box<dyn Collector<E>>, f: F) -> Self {
        Self { collector, f }
    }
}

#[async_trait]
impl<E1, E2, F> Collector<E2> for CollectorMap<E1, F>
where
    E1: Send + Sync + 'static,
    E2: Send + Sync + 'static,
    F: Fn(E1) -> E2 + Send + Sync + Clone + 'static,
{
    async fn get_event_stream(&self) -> Result<CollectorStream<'_, E2>> {
        let stream = self.collector.get_event_stream().await?;
        let f = self.f.clone();
        let stream = stream.map(f);
        Ok(Box::pin(stream))
    }
}

/// A wrapper around an [Executor] that maps incoming actions to a different type.
///
/// This adapter allows you to transform actions from one type to another before
/// they are executed. This is useful when you need to:
/// - Convert between different action representations
/// - Filter out unwanted actions
/// - Modify actions before execution
///
/// # Type Parameters
///
/// * `A` - The original action type
/// * `F` - The function type that maps actions
pub struct ExecutorMap<A, F> {
    executor: Box<dyn Executor<A>>,
    f: F,
}

impl<A, F> ExecutorMap<A, F> {
    /// Creates a new executor map with the given executor and mapping function.
    pub fn new(executor: Box<dyn Executor<A>>, f: F) -> Self {
        Self { executor, f }
    }
}

#[async_trait]
impl<A1, A2, F> Executor<A1> for ExecutorMap<A2, F>
where
    A1: Send + Sync + 'static,
    A2: Send + Sync + 'static,
    F: Fn(A1) -> Option<A2> + Send + Sync + Clone + 'static,
{
    async fn execute(&self, action: A1) -> Result<()> {
        let action = (self.f)(action);
        match action {
            Some(action) => self.executor.execute(action).await,
            None => Ok(()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use tokio::sync::Mutex;
    use tokio_stream::StreamExt;

    // Test event and action types
    #[derive(Debug, Clone, PartialEq)]
    struct TestEvent(u64);

    #[derive(Debug, Clone, PartialEq)]
    struct TestAction(String);

    // Mock collector implementation
    struct MockCollector {
        events: Vec<TestEvent>,
    }

    #[async_trait]
    impl Collector<TestEvent> for MockCollector {
        async fn get_event_stream(&self) -> Result<CollectorStream<'_, TestEvent>> {
            let events = self.events.clone();
            Ok(Box::pin(tokio_stream::iter(events)))
        }
    }

    // Mock strategy implementation
    struct MockStrategy {
        state: Arc<Mutex<u64>>,
    }

    #[async_trait]
    impl Strategy<TestEvent, TestAction> for MockStrategy {
        async fn sync_state(&mut self) -> Result<()> {
            let mut state = self.state.lock().await;
            *state = 0;
            Ok(())
        }

        async fn process_event(&mut self, event: TestEvent) -> Vec<TestAction> {
            let mut state = self.state.lock().await;
            *state += event.0;
            vec![TestAction(format!("processed_{}", event.0))]
        }
    }

    // Mock executor implementation
    struct MockExecutor {
        executed_actions: Arc<Mutex<Vec<TestAction>>>,
    }

    #[async_trait]
    impl Executor<TestAction> for MockExecutor {
        async fn execute(&self, action: TestAction) -> Result<()> {
            let mut actions = self.executed_actions.lock().await;
            actions.push(action);
            Ok(())
        }
    }

    #[tokio::test]
    async fn test_collector() {
        let collector = MockCollector {
            events: vec![TestEvent(1), TestEvent(2), TestEvent(3)],
        };

        let mut stream = collector.get_event_stream().await.unwrap();
        let mut events = Vec::new();
        while let Some(event) = stream.next().await {
            events.push(event);
        }

        assert_eq!(events, vec![TestEvent(1), TestEvent(2), TestEvent(3)]);
    }

    #[tokio::test]
    async fn test_strategy() {
        let state = Arc::new(Mutex::new(0));
        let mut strategy = MockStrategy {
            state: Arc::clone(&state),
        };

        // Test state synchronization
        strategy.sync_state().await.unwrap();
        assert_eq!(*state.lock().await, 0);

        // Test event processing
        let actions = strategy.process_event(TestEvent(5)).await;
        assert_eq!(actions, vec![TestAction("processed_5".to_string())]);
        assert_eq!(*state.lock().await, 5);
    }

    #[tokio::test]
    async fn test_executor() {
        let executed_actions = Arc::new(Mutex::new(Vec::new()));
        let executor = MockExecutor {
            executed_actions: Arc::clone(&executed_actions),
        };

        let action = TestAction("test_action".to_string());
        executor.execute(action.clone()).await.unwrap();

        let actions = executed_actions.lock().await;
        assert_eq!(*actions, vec![action]);
    }

    #[tokio::test]
    async fn test_collector_map() {
        let collector = MockCollector {
            events: vec![TestEvent(1), TestEvent(2)],
        };

        let mapped_collector = CollectorMap::new(Box::new(collector), |event: TestEvent| {
            TestEvent(event.0 * 2)
        });

        let mut stream = mapped_collector.get_event_stream().await.unwrap();
        let mut events = Vec::new();
        while let Some(event) = stream.next().await {
            events.push(event);
        }

        assert_eq!(events, vec![TestEvent(2), TestEvent(4)]);
    }

    #[tokio::test]
    async fn test_executor_map() {
        let executed_actions = Arc::new(Mutex::new(Vec::new()));
        let executor = MockExecutor {
            executed_actions: Arc::clone(&executed_actions),
        };

        let mapped_executor = ExecutorMap::new(Box::new(executor), |action: TestAction| {
            Some(TestAction(format!("mapped_{}", action.0)))
        });

        let action = TestAction("test".to_string());
        mapped_executor.execute(action).await.unwrap();

        let actions = executed_actions.lock().await;
        assert_eq!(*actions, vec![TestAction("mapped_test".to_string())]);
    }
}