foxhunt/market-data/src/orderbook.rs

//! # Order Book Module
//!
//! This module provides repository abstractions and implementations for storing,
//! retrieving, and managing order book data. It handles order book snapshots,
//! bid/ask levels, and liquidity analysis for trading instruments.
//!
//! ## Features
//!
//! - Storage and retrieval of order book levels
//! - Complete order book snapshot management
//! - Historical order book data queries
//! - Best bid/ask price retrieval
//! - Liquidity profile analysis
//! - Data cleanup and maintenance operations
//!
//! ## Usage
//!
//! ```rust
//! use market_data::orderbook::{OrderBookRepository, PostgresOrderBookRepository};
//! use market_data::models::OrderBook;
//! use sqlx::PgPool;
//!
//! # async fn example(pool: PgPool) -> Result<(), Box<dyn std::error::Error>> {
//! let repo = PostgresOrderBookRepository::new(pool);
//!
//! // Get latest order book for a symbol
//! let order_book = repo.get_latest_order_book("AAPL").await?;
//!
//! // Get best bid and ask
//! let best_prices = repo.get_best_bid_ask("AAPL").await?;
//! # Ok(())
//! # }
//! ```

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use sqlx::{PgPool, Row};
use std::collections::HashMap;

use crate::{
    error::{MarketDataError, MarketDataResult},
    models::{BookSide, OrderBook, OrderBookLevelDb},
};

/// Repository trait for order book data operations
///
/// This trait defines the interface for storing, retrieving, and managing
/// order book data. Implementations should provide efficient access patterns
/// optimized for real-time trading and historical analysis.
///
/// The trait supports:
/// - Individual level and complete order book storage
///
/// - Real-time best bid/ask queries
/// - Historical order book reconstruction
///
/// - Liquidity analysis and profiling
/// - Data maintenance and cleanup
#[async_trait]
pub trait OrderBookRepository {
    /// Store order book levels for a symbol
    ///
    /// Stores multiple order book levels in a single batch operation.
    ///
    /// This is optimized for storing complete order book snapshots efficiently.
    ///
    /// # Arguments
    ///
    /// * `levels` - Slice of order book levels to store
    ///
    /// # Returns
    ///
    /// `Ok(())` on success, or a `MarketDataError` if the operation fails
    ///
    /// # Errors
    ///
    /// - `MarketDataError::InvalidSymbol` if any symbol is invalid
    /// - `MarketDataError::Database` if the database transaction fails
    async fn store_order_book_levels(&self, levels: &[OrderBookLevelDb]) -> MarketDataResult<()>;

    /// Store a complete order book snapshot
    ///
    /// Stores a complete order book with all bid and ask levels.
    ///
    /// This is a convenience method that extracts levels from the order book
    /// and stores them using `store_order_book_levels`.
    ///
    /// # Arguments
    ///
    /// * `order_book` - Complete order book snapshot to store
    ///
    /// # Returns
    ///
    /// `Ok(())` on success, or a `MarketDataError` if the operation fails
    ///
    /// # Errors
    ///
    /// - `MarketDataError::InvalidSymbol` if the symbol is invalid
    /// - `MarketDataError::Database` if the database operation fails
    async fn store_order_book(&self, order_book: &OrderBook) -> MarketDataResult<()>;

    /// Get the latest order book for a symbol
    ///
    /// Retrieves the most recent complete order book snapshot for the specified symbol.
    ///
    /// The returned order book includes all available bid and ask levels sorted appropriately.
    ///
    /// # Arguments
    ///
    /// * `symbol` - Trading symbol to query
    ///
    /// # Returns
    ///
    /// `Some(order_book)` if found, `None` if no data exists, or a `MarketDataError`
    ///
    /// # Errors
    ///
    /// - `MarketDataError::InvalidSymbol` if the symbol is invalid
    /// - `MarketDataError::Database` if the query fails
    async fn get_latest_order_book(&self, symbol: &str) -> MarketDataResult<Option<OrderBook>>;

    /// Get order book levels for a symbol at a specific time
    ///
    /// Retrieves order book levels for a specific timestamp, optionally
    /// limiting the number of levels returned. Levels are sorted by price.
    ///
    /// # Arguments
    ///
    /// * `symbol` - Trading symbol to query
    /// * `timestamp` - Specific timestamp to query
    ///
    /// * `max_levels` - Optional limit on number of levels (defaults to 50)
    ///
    /// # Returns
    ///
    /// Vector of order book levels sorted by price, or a `MarketDataError`
    ///
    /// # Errors
    ///
    /// - `MarketDataError::InvalidSymbol` if the symbol is invalid
    /// - `MarketDataError::Database` if the query fails
    async fn get_order_book_levels(
        &self,
        symbol: &str,
        timestamp: DateTime<Utc>,
        max_levels: Option<i32>,
    ) -> MarketDataResult<Vec<OrderBookLevelDb>>;

    /// Get order book history for a time range
    ///
    /// Retrieves historical order book snapshots within the specified time range.
    ///
    /// Each snapshot represents the complete order book state at a specific time.
    ///
    /// # Arguments
    ///
    /// * `symbol` - Trading symbol to query
    /// * `from` - Start of time range (inclusive)
    ///
    /// * `to` - End of time range (inclusive)
    ///
    /// # Returns
    ///
    /// Vector of order books ordered by timestamp, or a `MarketDataError`
    ///
    /// # Errors
    ///
    /// - `MarketDataError::InvalidSymbol` if the symbol is invalid
    /// - `MarketDataError::InvalidTimeRange` if from >= to
    ///
    /// - `MarketDataError::Database` if the query fails
    async fn get_order_book_history(
        &self,
        symbol: &str,
        from: DateTime<Utc>,
        to: DateTime<Utc>,
    ) -> MarketDataResult<Vec<OrderBook>>;

    /// Get the best bid and ask for a symbol
    ///
    /// Retrieves the highest bid price and lowest ask price for the specified symbol.
    ///
    /// This represents the current market spread and best available prices.
    ///
    /// # Arguments
    ///
    /// * `symbol` - Trading symbol to query
    ///
    /// # Returns
    ///
    /// `Some((best_bid, best_ask))` if both sides exist, `None` otherwise, or a `MarketDataError`
    ///
    /// # Errors
    ///
    /// - `MarketDataError::InvalidSymbol` if the symbol is invalid
    /// - `MarketDataError::Database` if the query fails
    async fn get_best_bid_ask(
        &self,
        symbol: &str,
    ) -> MarketDataResult<Option<(OrderBookLevelDb, OrderBookLevelDb)>>;

    /// Get aggregated liquidity at price levels
    ///
    /// Retrieves the complete liquidity profile showing all bid and ask levels
    /// at a specific timestamp. Useful for market depth analysis.
    ///
    /// # Arguments
    ///
    /// * `symbol` - Trading symbol to analyze
    /// * `timestamp` - Specific timestamp to query
    ///
    /// # Returns
    ///
    /// HashMap with bid and ask sides mapped to their respective levels, or a `MarketDataError`
    ///
    /// # Errors
    ///
    /// - `MarketDataError::InvalidSymbol` if the symbol is invalid
    /// - `MarketDataError::Database` if the query fails
    async fn get_liquidity_profile(
        &self,
        symbol: &str,
        timestamp: DateTime<Utc>,
    ) -> MarketDataResult<HashMap<BookSide, Vec<OrderBookLevelDb>>>;

    /// Delete old order book data before a given timestamp
    ///
    /// Removes historical order book data older than the specified timestamp.
    ///
    /// This is useful for data retention management and storage optimization.
    ///
    /// # Arguments
    ///
    /// * `before` - Timestamp before which all data will be deleted
    ///
    /// # Returns
    ///
    /// Number of records deleted, or a `MarketDataError`
    ///
    /// # Errors
    ///
    /// - `MarketDataError::Database` if the deletion fails
    async fn cleanup_old_order_book_data(&self, before: DateTime<Utc>) -> MarketDataResult<u64>;
}

/// PostgreSQL implementation of OrderBookRepository
///
/// Provides a production-ready implementation of the `OrderBookRepository` trait
/// using PostgreSQL as the backend storage. This implementation is optimized
/// for high-frequency order book updates and fast retrieval patterns.
///
/// ## Features
///
/// - Transactional batch operations for atomic updates
/// - Optimized queries for time-series data
///
/// - Proper bid/ask sorting and level organization
/// - Input validation and error handling
///
/// - Conflict resolution with upsert semantics
pub struct PostgresOrderBookRepository {
    /// PostgreSQL connection pool for database operations
    pool: PgPool,
}

impl PostgresOrderBookRepository {
    /// Create a new PostgreSQL order book repository
    ///
    /// # Arguments
    ///
    /// * `pool` - PostgreSQL connection pool
    ///
    /// # Returns
    ///
    /// A new `PostgresOrderBookRepository` instance
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }

    /// Validate that a symbol meets format requirements
    ///
    /// Ensures the symbol is non-empty and within length limits.
    ///
    /// # Arguments
    ///
    /// * `symbol` - Symbol to validate
    ///
    /// # Returns
    ///
    /// `Ok(())` if valid, `MarketDataError::InvalidSymbol` otherwise
    async fn validate_symbol(&self, symbol: &str) -> MarketDataResult<()> {
        if symbol.is_empty() || symbol.len() > 20 {
            return Err(MarketDataError::InvalidSymbol {
                symbol: symbol.to_string(),
            });
        }
        Ok(())
    }

    /// Validate that a time range is logically correct
    ///
    /// Ensures the start time is before the end time.
    ///
    /// # Arguments
    ///
    /// * `from` - Start time
    /// * `to` - End time
    ///
    /// # Returns
    ///
    /// `Ok(())` if valid, `MarketDataError::InvalidTimeRange` otherwise
    async fn validate_time_range(
        &self,
        from: DateTime<Utc>,
        to: DateTime<Utc>,
    ) -> MarketDataResult<()> {
        if from >= to {
            return Err(MarketDataError::InvalidTimeRange { from, to });
        }
        Ok(())
    }
}

#[async_trait]
impl OrderBookRepository for PostgresOrderBookRepository {
    async fn store_order_book_levels(&self, levels: &[OrderBookLevelDb]) -> MarketDataResult<()> {
        if levels.is_empty() {
            return Ok(());
        }

        // Validate all symbols first
        for level in levels {
            self.validate_symbol(&level.symbol).await?;
        }

        let mut tx = self.pool.begin().await?;

        for level in levels {
            sqlx::query(
                r#"
                INSERT INTO order_book_levels (id, symbol, timestamp, side, price, quantity, level, created_at)
                VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                ON CONFLICT (symbol, timestamp, side, level) DO UPDATE SET
                    price = EXCLUDED.price,
                    quantity = EXCLUDED.quantity
                "#
            )
            .bind(level.id)
            .bind(&level.symbol)
            .bind(level.timestamp)
            .bind(level.side as BookSide)
            .bind(level.price)
            .bind(level.quantity)
            .bind(level.level)
            .bind(level.created_at)
            .execute(&mut *tx)
            .await?;
        }

        tx.commit().await?;
        Ok(())
    }

    async fn store_order_book(&self, order_book: &OrderBook) -> MarketDataResult<()> {
        self.validate_symbol(&order_book.symbol).await?;

        let mut levels = Vec::new();
        levels.extend(order_book.bids.clone());
        levels.extend(order_book.asks.clone());

        self.store_order_book_levels(&levels).await
    }

    async fn get_latest_order_book(&self, symbol: &str) -> MarketDataResult<Option<OrderBook>> {
        self.validate_symbol(symbol).await?;

        // Get the latest timestamp for this symbol
        let latest_timestamp = sqlx::query(
            "SELECT MAX(timestamp) as latest_timestamp FROM order_book_levels WHERE symbol = $1",
        )
        .bind(symbol)
        .fetch_one(&self.pool)
        .await?;

        if let Some(timestamp) =
            latest_timestamp.get::<Option<DateTime<Utc>>, _>("latest_timestamp")
        {
            let levels = self.get_order_book_levels(symbol, timestamp, None).await?;

            if levels.is_empty() {
                return Ok(None);
            }

            let mut order_book = OrderBook::new(symbol.to_string(), timestamp);

            for level in levels {
                match level.side {
                    BookSide::Bid => order_book.bids.push(level),
                    BookSide::Ask => order_book.asks.push(level),
                }
            }

            // Sort bids by price descending (highest first)
            order_book.bids.sort_by(|a, b| b.price.cmp(&a.price));
            // Sort asks by price ascending (lowest first)
            order_book.asks.sort_by(|a, b| a.price.cmp(&b.price));

            Ok(Some(order_book))
        } else {
            Ok(None)
        }
    }

    async fn get_order_book_levels(
        &self,
        symbol: &str,
        timestamp: DateTime<Utc>,
        max_levels: Option<i32>,
    ) -> MarketDataResult<Vec<OrderBookLevelDb>> {
        self.validate_symbol(symbol).await?;

        let limit = max_levels.unwrap_or(50); // Default to 50 levels

        let rows = sqlx::query(
            r#"
            SELECT id, symbol, timestamp, side, price, quantity, level, created_at
            FROM order_book_levels
            WHERE symbol = $1 AND timestamp = $2
            ORDER BY 
                CASE WHEN side = 'bid' THEN -price ELSE price END,
                level
            LIMIT $3
            "#,
        )
        .bind(symbol)
        .bind(timestamp)
        .bind(limit as i64)
        .fetch_all(&self.pool)
        .await?;

        let levels = rows
            .into_iter()
            .map(|row| OrderBookLevelDb {
                id: row.get("id"),
                symbol: row.get("symbol"),
                timestamp: row.get("timestamp"),
                side: row.get("side"),
                price: row.get("price"),
                quantity: row.get("quantity"),
                level: row.get("level"),
                created_at: row.get("created_at"),
            })
            .collect();

        Ok(levels)
    }

    async fn get_order_book_history(
        &self,
        symbol: &str,
        from: DateTime<Utc>,
        to: DateTime<Utc>,
    ) -> MarketDataResult<Vec<OrderBook>> {
        self.validate_symbol(symbol).await?;
        self.validate_time_range(from, to).await?;

        // Get distinct timestamps in the range
        let timestamps = sqlx::query(
            r#"
            SELECT DISTINCT timestamp
            FROM order_book_levels
            WHERE symbol = $1 AND timestamp >= $2 AND timestamp <= $3
            ORDER BY timestamp ASC
            "#,
        )
        .bind(symbol)
        .bind(from)
        .bind(to)
        .fetch_all(&self.pool)
        .await?;

        let mut order_books = Vec::new();

        for timestamp_row in timestamps {
            let timestamp: DateTime<Utc> = timestamp_row.get("timestamp");
            let levels = self.get_order_book_levels(symbol, timestamp, None).await?;

            if !levels.is_empty() {
                let mut order_book = OrderBook::new(symbol.to_string(), timestamp);

                for level in levels {
                    match level.side {
                        BookSide::Bid => order_book.bids.push(level),
                        BookSide::Ask => order_book.asks.push(level),
                    }
                }

                // Sort bids by price descending, asks by price ascending
                order_book.bids.sort_by(|a, b| b.price.cmp(&a.price));
                order_book.asks.sort_by(|a, b| a.price.cmp(&b.price));

                order_books.push(order_book);
            }
        }

        Ok(order_books)
    }

    async fn get_best_bid_ask(
        &self,
        symbol: &str,
    ) -> MarketDataResult<Option<(OrderBookLevelDb, OrderBookLevelDb)>> {
        self.validate_symbol(symbol).await?;

        let best_bid = sqlx::query(
            r#"
            SELECT id, symbol, timestamp, side, price, quantity, level, created_at
            FROM order_book_levels
            WHERE symbol = $1 AND side = 'bid'
            ORDER BY timestamp DESC, price DESC
            LIMIT 1
            "#,
        )
        .bind(symbol)
        .fetch_optional(&self.pool)
        .await?;

        let best_ask = sqlx::query(
            r#"
            SELECT id, symbol, timestamp, side, price, quantity, level, created_at
            FROM order_book_levels
            WHERE symbol = $1 AND side = 'ask'
            ORDER BY timestamp DESC, price ASC
            LIMIT 1
            "#,
        )
        .bind(symbol)
        .fetch_optional(&self.pool)
        .await?;

        match (best_bid, best_ask) {
            (Some(bid_row), Some(ask_row)) => {
                let bid_level = OrderBookLevelDb {
                    id: bid_row.get("id"),
                    symbol: bid_row.get("symbol"),
                    timestamp: bid_row.get("timestamp"),
                    side: bid_row.get("side"),
                    price: bid_row.get("price"),
                    quantity: bid_row.get("quantity"),
                    level: bid_row.get("level"),
                    created_at: bid_row.get("created_at"),
                };

                let ask_level = OrderBookLevelDb {
                    id: ask_row.get("id"),
                    symbol: ask_row.get("symbol"),
                    timestamp: ask_row.get("timestamp"),
                    side: ask_row.get("side"),
                    price: ask_row.get("price"),
                    quantity: ask_row.get("quantity"),
                    level: ask_row.get("level"),
                    created_at: ask_row.get("created_at"),
                };

                Ok(Some((bid_level, ask_level)))
            },
            _ => Ok(None),
        }
    }

    async fn get_liquidity_profile(
        &self,
        symbol: &str,
        timestamp: DateTime<Utc>,
    ) -> MarketDataResult<HashMap<BookSide, Vec<OrderBookLevelDb>>> {
        self.validate_symbol(symbol).await?;

        let levels = self.get_order_book_levels(symbol, timestamp, None).await?;

        let mut profile = HashMap::new();
        profile.insert(BookSide::Bid, Vec::new());
        profile.insert(BookSide::Ask, Vec::new());

        for level in levels {
            profile
                .entry(level.side)
                .or_insert_with(Vec::new)
                .push(level);
        }

        // Sort by price (descending for bids, ascending for asks)
        if let Some(bids) = profile.get_mut(&BookSide::Bid) {
            bids.sort_by(|a, b| b.price.cmp(&a.price));
        }
        if let Some(asks) = profile.get_mut(&BookSide::Ask) {
            asks.sort_by(|a, b| a.price.cmp(&b.price));
        }

        Ok(profile)
    }

    async fn cleanup_old_order_book_data(&self, before: DateTime<Utc>) -> MarketDataResult<u64> {
        let result = sqlx::query("DELETE FROM order_book_levels WHERE timestamp < $1")
            .bind(before)
            .execute(&self.pool)
            .await?;

        Ok(result.rows_affected())
    }
}