atelier_data 0.0.15

use crate::{
    errors::OrderbookError,
    orderbooks::{OrderbookUpdate, OrderbookUpdateType},
    utils::current_timestamp_ms,
};
use rust_decimal::Decimal;
use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, str::FromStr};

// ─────────────────────────────────────────────────────────────────────────────
// NormalizedDelta — exchange-agnostic orderbook update
// ─────────────────────────────────────────────────────────────────────────────

/// Exchange-agnostic orderbook delta/snapshot representation.
///
/// Every exchange decoder converts its raw response into a `NormalizedDelta`
/// before feeding it to [`OrderbookDelta::process`]. This keeps all
/// exchange-specific parsing confined to the decoder layer.
#[derive(Debug, Clone)]
pub struct NormalizedDelta {
    /// Trading pair symbol (e.g. `"BTCUSDT"`, `"BTC-USD"`).
    pub symbol: String,
    /// Bid levels as `(price_str, size_str)` pairs.
    pub bids: Vec<(String, String)>,
    /// Ask levels as `(price_str, size_str)` pairs.
    pub asks: Vec<(String, String)>,
    /// Sequence / update ID for gap detection.
    pub update_id: u64,
    /// Cross-sequence number (0 if the exchange does not provide one).
    pub sequence: u64,
    /// `true` → full book reset; `false` → incremental update.
    pub is_snapshot: bool,
}

/// A Delta Orderbook state from WebSocket updates
///
/// Follows these rules:
/// - Snapshot: reset entire book
/// - Delta with size=0: delete the price level
/// - Delta with new price: insert the level
/// - Delta with existing price: update the size
#[derive(Debug, Clone, Serialize)]
pub struct OrderbookDelta {
    /// Bid side: price -> size (BTreeMap sorts ascending, use .last() for best bid)
    pub bids: BTreeMap<Decimal, Decimal>,
    /// Ask side: price -> size (BTreeMap sorts ascending, use .first() for best ask)
    pub asks: BTreeMap<Decimal, Decimal>,
    /// Last update ID for sequencing validation
    pub last_update_id: u64,
    /// Sequence number
    pub sequence: u64,
    /// Symbol this book is tracking
    pub symbol: String,
    /// Exchange from where the data is tracked
    pub exchange: String,
    /// Whether if an initial snapshot was received
    pub initialized: bool,
    /// Count of updates applied since last snapshot
    pub delta_count: u64,
}

impl OrderbookDelta {
    /// Create a new orderbook delta for a symbol
    pub fn new(symbol: impl Into<String>) -> Self {
        Self {
            bids: BTreeMap::new(),
            asks: BTreeMap::new(),
            last_update_id: 0,
            sequence: 0,
            symbol: symbol.into(),
            exchange: "".to_string(),
            initialized: false,
            delta_count: 0,
        }
    }

    /// Check if the orderbook has been initialized with a snapshot
    #[inline]
    pub fn is_initialized(&self) -> bool {
        self.initialized
    }

    /// Get the symbol this orderbook is tracking
    #[inline]
    pub fn symbol(&self) -> &str {
        &self.symbol
    }

    /// Get the symbol this orderbook is tracking
    #[inline]
    pub fn exchange(&self) -> &str {
        &self.exchange
    }

    /// Get the last update ID
    #[inline]
    pub fn last_update_id(&self) -> u64 {
        self.last_update_id
    }

    /// Get the sequence number
    #[inline]
    pub fn sequence(&self) -> u64 {
        self.sequence
    }

    /// Get count of deltas applied since last snapshot
    #[inline]
    pub fn delta_count(&self) -> u64 {
        self.delta_count
    }

    /// Process an exchange-agnostic orderbook update (snapshot or delta).
    ///
    /// Each exchange decoder is responsible for converting its raw response
    /// into a [`NormalizedDelta`] before calling this method.
    pub fn process(
        &mut self,
        delta: &NormalizedDelta,
    ) -> Result<OrderbookUpdate, OrderbookError> {
        // Validate symbol
        if !self.symbol.is_empty() && delta.symbol != self.symbol {
            return Err(OrderbookError::SymbolMismatch {
                expected: self.symbol.clone(),
                received: delta.symbol.clone(),
            });
        }

        if delta.is_snapshot {
            self.apply_snapshot(delta)
        } else {
            self.apply_delta(delta)
        }
    }

    /// Apply a snapshot (full reset)
    fn apply_snapshot(
        &mut self,
        delta: &NormalizedDelta,
    ) -> Result<OrderbookUpdate, OrderbookError> {
        self.bids.clear();
        self.asks.clear();

        let mut bids_modified = 0;
        let mut asks_modified = 0;

        for (price_str, size_str) in &delta.bids {
            let price = Self::parse_decimal(price_str)?;
            let size = Self::parse_decimal(size_str)?;
            if size > Decimal::ZERO {
                self.bids.insert(price, size);
                bids_modified += 1;
            }
        }

        for (price_str, size_str) in &delta.asks {
            let price = Self::parse_decimal(price_str)?;
            let size = Self::parse_decimal(size_str)?;
            if size > Decimal::ZERO {
                self.asks.insert(price, size);
                asks_modified += 1;
            }
        }

        self.last_update_id = delta.update_id;
        self.sequence = delta.sequence;
        self.symbol.clone_from(&delta.symbol);
        self.initialized = true;
        self.delta_count = 0;

        Ok(OrderbookUpdate {
            update_type: OrderbookUpdateType::Snapshot,
            bids_modified,
            asks_modified,
            levels_deleted: 0,
            levels_inserted: bids_modified + asks_modified,
            was_reset: true,
        })
    }

    /// Apply a delta update
    fn apply_delta(
        &mut self,
        delta: &NormalizedDelta,
    ) -> Result<OrderbookUpdate, OrderbookError> {
        if !self.initialized {
            return Err(OrderbookError::NotInitialized);
        }

        let mut bids_modified = 0;
        let mut asks_modified = 0;
        let mut levels_deleted = 0;
        let mut levels_inserted = 0;

        for (price_str, size_str) in &delta.bids {
            let price = Self::parse_decimal(price_str)?;
            let size = Self::parse_decimal(size_str)?;
            let (deleted, inserted) = Self::apply_level(&mut self.bids, price, size);
            bids_modified += 1;
            if deleted {
                levels_deleted += 1;
            }
            if inserted {
                levels_inserted += 1;
            }
        }

        for (price_str, size_str) in &delta.asks {
            let price = Self::parse_decimal(price_str)?;
            let size = Self::parse_decimal(size_str)?;
            let (deleted, inserted) = Self::apply_level(&mut self.asks, price, size);
            asks_modified += 1;
            if deleted {
                levels_deleted += 1;
            }
            if inserted {
                levels_inserted += 1;
            }
        }

        self.last_update_id = delta.update_id;
        self.sequence = delta.sequence;
        self.delta_count += 1;

        Ok(OrderbookUpdate {
            update_type: OrderbookUpdateType::Delta,
            bids_modified,
            asks_modified,
            levels_deleted,
            levels_inserted,
            was_reset: false,
        })
    }

    /// Produces an [`OrderbookSnapshot`] from the current delta state.
    ///
    /// This is an internal helper that extracts all relevant data from an
    /// [`OrderbookDelta`] into a serializable snapshot structure.
    ///
    /// # Arguments
    ///
    /// * `ob` - Mutable Reference to Self `OrderbookDelta` to generate the snapshot from
    ///
    /// # Returns
    ///
    /// An [`OrderbookSnapshot`] containing:
    /// - Current timestamp (capture time, not exchange time)
    /// - All price levels as string tuples to preserve decimal precision
    /// - Derived metrics rounded to appropriate decimal places
    pub fn produce_snapshot(ob: &mut Self) -> OrderbookSnapshot {
        let bids: Vec<[String; 2]> = ob
            .top_bids(ob.bid_depth())
            .iter()
            .map(|(p, s)| [p.to_string(), s.to_string()])
            .collect();

        let asks: Vec<[String; 2]> = ob
            .top_asks(ob.ask_depth())
            .iter()
            .map(|(p, s)| [p.to_string(), s.to_string()])
            .collect();

        OrderbookSnapshot {
            timestamp_ms: current_timestamp_ms(),
            symbol: ob.symbol().to_string(),
            update_id: ob.last_update_id(),
            sequence: ob.sequence(),
            delta_count: ob.delta_count(),
            bid_depth: ob.bid_depth(),
            ask_depth: ob.ask_depth(),
            mid_price: ob.mid_price().map(|d| d.round_dp(8).to_string()),
            spread: ob.spread().map(|d| d.round_dp(8).to_string()),
            spread_bps: ob.spread_bps().map(|d| d.round_dp(4).to_string()),
            volume_imbalance: ob.volume_imbalance().map(|d| d.round_dp(6).to_string()),
            total_bid_volume: ob.total_bid_volume().round_dp(8).to_string(),
            total_ask_volume: ob.total_ask_volume().round_dp(8).to_string(),
            bids,
            asks,
        }
    }

    /// Apply a single level update to a side of the book
    /// Returns (was_deleted, was_inserted)
    fn apply_level(
        book_side: &mut BTreeMap<Decimal, Decimal>,
        price: Decimal,
        size: Decimal,
    ) -> (bool, bool) {
        if size == Decimal::ZERO {
            // Delete the entry (Level was dropped)
            let existed = book_side.remove(&price).is_some();
            (existed, false)
        } else {
            match book_side.entry(price) {
                std::collections::btree_map::Entry::Occupied(mut e) => {
                    // Update existing entry (Level was modified)
                    e.insert(size);
                    (false, false)
                }
                std::collections::btree_map::Entry::Vacant(e) => {
                    // Insert new entry (Level was created)
                    e.insert(size);
                    (false, true)
                }
            }
        }
    }

    /// Parse a string to Decimal
    fn parse_decimal(s: &str) -> Result<Decimal, OrderbookError> {
        Decimal::from_str(s).map_err(|e| OrderbookError::ParseError(e.to_string()))
    }

    // ═══════════════════════════════════════════════════════════════════════════
    // Book State Accessors
    // ═══════════════════════════════════════════════════════════════════════════

    /// Get the best bid (highest bid price)
    #[inline]
    pub fn best_bid(&self) -> Option<(Decimal, Decimal)> {
        self.bids.last_key_value().map(|(&p, &s)| (p, s))
    }

    /// Get the best ask (lowest ask price)
    #[inline]
    pub fn best_ask(&self) -> Option<(Decimal, Decimal)> {
        self.asks.first_key_value().map(|(&p, &s)| (p, s))
    }

    /// Get best bid/ask as (bid_price, bid_size, ask_price, ask_size)
    pub fn bbo(&self) -> Option<(Decimal, Decimal, Decimal, Decimal)> {
        match (self.best_bid(), self.best_ask()) {
            (Some((bp, bs)), Some((ap, az))) => Some((bp, bs, ap, az)),
            _ => None,
        }
    }

    /// Get mid price
    pub fn mid_price(&self) -> Option<Decimal> {
        match (self.best_bid(), self.best_ask()) {
            (Some((bid, _)), Some((ask, _))) => Some((bid + ask) / Decimal::TWO),
            _ => None,
        }
    }

    /// Get spread
    pub fn spread(&self) -> Option<Decimal> {
        match (self.best_bid(), self.best_ask()) {
            (Some((bid, _)), Some((ask, _))) => Some(ask - bid),
            _ => None,
        }
    }

    /// Get spread in basis points
    pub fn spread_bps(&self) -> Option<Decimal> {
        let spread = self.spread()?;
        let mid = self.mid_price()?;
        if mid > Decimal::ZERO {
            Some((spread / mid) * Decimal::from(10_000))
        } else {
            None
        }
    }

    /// Get number of bid levels
    #[inline]
    pub fn bid_depth(&self) -> usize {
        self.bids.len()
    }

    /// Get number of ask levels
    #[inline]
    pub fn ask_depth(&self) -> usize {
        self.asks.len()
    }

    /// Get total bid volume
    pub fn total_bid_volume(&self) -> Decimal {
        self.bids.values().copied().sum()
    }

    /// Get total ask volume
    pub fn total_ask_volume(&self) -> Decimal {
        self.asks.values().copied().sum()
    }

    /// Volume imbalance: (bid - ask) / (bid + ask) ∈ [-1, 1]
    pub fn volume_imbalance(&self) -> Option<Decimal> {
        let bid_vol = self.total_bid_volume();
        let ask_vol = self.total_ask_volume();
        let total = bid_vol + ask_vol;
        if total > Decimal::ZERO {
            Some((bid_vol - ask_vol) / total)
        } else {
            None
        }
    }

    /// Get top N bid levels (highest to lowest price)
    pub fn top_bids(&self, n: usize) -> Vec<(Decimal, Decimal)> {
        self.bids
            .iter()
            .rev()
            .take(n)
            .map(|(&p, &s)| (p, s))
            .collect()
    }

    /// Get top N ask levels (lowest to highest price)
    pub fn top_asks(&self, n: usize) -> Vec<(Decimal, Decimal)> {
        self.asks.iter().take(n).map(|(&p, &s)| (p, s)).collect()
    }

    /// Get bid volume within a price range from best bid
    pub fn bid_volume_within(&self, depth: Decimal) -> Decimal {
        let Some((best_bid, _)) = self.best_bid() else {
            return Decimal::ZERO;
        };
        let threshold = best_bid - depth;
        self.bids.range(threshold..).map(|(_, &size)| size).sum()
    }

    /// Get ask volume within a price range from best ask
    pub fn ask_volume_within(&self, depth: Decimal) -> Decimal {
        let Some((best_ask, _)) = self.best_ask() else {
            return Decimal::ZERO;
        };
        let threshold = best_ask + depth;
        self.asks.range(..=threshold).map(|(_, &size)| size).sum()
    }

    /// Create an OrderbookDelta from pre-built BTreeMaps (minimal metadata)
    ///
    /// Used when loading from CSV or Parquet where only price/size data is available.
    pub fn from_maps(
        symbol: impl Into<String>,
        exchange: impl Into<String>,
        bids: BTreeMap<Decimal, Decimal>,
        asks: BTreeMap<Decimal, Decimal>,
    ) -> Self {
        let bid_count = bids.len();
        let ask_count = asks.len();

        Self {
            symbol: symbol.into(),
            exchange: exchange.into(),
            bids,
            asks,
            last_update_id: 0,
            sequence: 0,
            initialized: bid_count > 0 || ask_count > 0,
            delta_count: 0,
        }
    }

    /// Create an OrderbookDelta from a full snapshot with metadata
    ///
    /// Used when loading from JSON where all metadata is preserved.
    pub fn from_snapshot(
        symbol: impl Into<String>,
        exchange: impl Into<String>,
        bids: BTreeMap<Decimal, Decimal>,
        asks: BTreeMap<Decimal, Decimal>,
        update_id: u64,
        sequence: u64,
        delta_count: u64,
    ) -> Self {
        let bid_count = bids.len();
        let ask_count = asks.len();

        Self {
            symbol: symbol.into(),
            exchange: exchange.into(),
            bids,
            asks,
            last_update_id: update_id,
            sequence,
            initialized: bid_count > 0 || ask_count > 0,
            delta_count,
        }
    }

    /// Create an OrderbookDelta from vectors of (price, size) tuples
    ///
    /// Useful for testing or manual construction.
    pub fn from_levels(
        symbol: impl Into<String>,
        exchange: impl Into<String>,
        bids: impl IntoIterator<Item = (Decimal, Decimal)>,
        asks: impl IntoIterator<Item = (Decimal, Decimal)>,
    ) -> Self {
        let bids: BTreeMap<Decimal, Decimal> = bids.into_iter().collect();
        let asks: BTreeMap<Decimal, Decimal> = asks.into_iter().collect();
        Self::from_maps(symbol, exchange, bids, asks)
    }
}

/// Full orderbook snapshot for JSON serialization
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct OrderbookSnapshot {
    pub timestamp_ms: u64,
    pub symbol: String,
    pub update_id: u64,
    pub sequence: u64,
    pub delta_count: u64,
    pub bid_depth: usize,
    pub ask_depth: usize,
    pub mid_price: Option<String>,
    pub spread: Option<String>,
    pub spread_bps: Option<String>,
    pub volume_imbalance: Option<String>,
    pub total_bid_volume: String,
    pub total_ask_volume: String,
    pub bids: Vec<[String; 2]>,
    pub asks: Vec<[String; 2]>,
}

impl OrderbookSnapshot {
    /// Pretty print the orderbook
    pub fn display(&self, levels: usize) {
        println!("╔══════════════════════════════════════════════════════════════╗");
        println!(
            "║  {} | update_id: {} | seq: {} | deltas: {}",
            self.symbol, self.update_id, self.sequence, self.delta_count
        );
        println!("╠══════════════════════════════════════════════════════════════╣");
        println!(
            "║  {:^15} │ {:^12} ║ {:^12} │ {:^15}  ║",
            "ASK PRICE", "ASK SIZE", "BID SIZE", "BID PRICE"
        );
        println!("╠══════════════════════════════════════════════════════════════╣");

        let display_levels = levels.min(self.bids.len()).min(self.asks.len());

        // Show asks in reverse (highest to lowest) then bids (highest to lowest)
        let asks_rev: Vec<_> = self.asks.iter().take(display_levels).collect();

        for (i, [ask_price, ask_size]) in asks_rev.iter().rev().enumerate() {
            if i < display_levels {
                let [bid_price, bid_size] = &self.bids[display_levels - 1 - i];
                println!(
                    "║  {:>15} │ {:>12} ║ {:>12} │ {:>15}  ║",
                    ask_price, ask_size, bid_size, bid_price
                );
            }
        }
        println!("╚══════════════════════════════════════════════════════════════╝");
    }
}