polynode 0.12.0

use std::collections::HashMap;
use std::time::{Duration, Instant};
use crate::types::orderbook::{OrderbookLevel, BookSnapshot, BookUpdate, PriceChange};

struct BookEntry {
    bids: Vec<OrderbookLevel>,
    asks: Vec<OrderbookLevel>,
    last_updated: Instant,
}

/// Local orderbook state manager. Apply snapshots and deltas to maintain
/// a sorted local copy of the book for any number of assets.
pub struct LocalOrderbook {
    books: HashMap<String, BookEntry>,
}

impl LocalOrderbook {
    pub fn new() -> Self {
        Self {
            books: HashMap::new(),
        }
    }

    /// Replace the full book for an asset.
    pub fn apply_snapshot(&mut self, snap: &BookSnapshot) {
        let mut bids = snap.bids.clone();
        let mut asks = snap.asks.clone();
        // Sort bids descending (best bid first), asks ascending (best ask first)
        bids.sort_by(|a, b| {
            let pa: f64 = b.price.parse().unwrap_or(0.0);
            let pb: f64 = a.price.parse().unwrap_or(0.0);
            pa.partial_cmp(&pb).unwrap_or(std::cmp::Ordering::Equal)
        });
        asks.sort_by(|a, b| {
            let pa: f64 = a.price.parse().unwrap_or(0.0);
            let pb: f64 = b.price.parse().unwrap_or(0.0);
            pa.partial_cmp(&pb).unwrap_or(std::cmp::Ordering::Equal)
        });
        self.books.insert(snap.asset_id.clone(), BookEntry {
            bids,
            asks,
            last_updated: Instant::now(),
        });
    }

    /// Apply an incremental delta. A level with size "0" means removal.
    pub fn apply_update(&mut self, update: &BookUpdate) {
        let entry = self.books.entry(update.asset_id.clone())
            .or_insert_with(|| BookEntry {
                bids: Vec::new(),
                asks: Vec::new(),
                last_updated: Instant::now(),
            });
        apply_deltas(&mut entry.bids, &update.bids, true);
        apply_deltas(&mut entry.asks, &update.asks, false);
        entry.last_updated = Instant::now();
    }

    /// Apply a `price_change` event — PM's primary incremental update.
    /// Each asset in the change carries (price, size, side); size="0" removes
    /// the level, otherwise it upserts. Requires server ≥ the 2026-04-19 fix
    /// that forwards size/side fields; older servers produce events with
    /// those fields absent and this is a no-op for them.
    pub fn apply_price_change(&mut self, change: &PriceChange) {
        let now = Instant::now();
        for a in &change.assets {
            let (Some(size), Some(side)) = (a.size.as_ref(), a.side.as_ref()) else { continue };
            let entry = self.books.entry(a.asset_id.clone())
                .or_insert_with(|| BookEntry {
                    bids: Vec::new(),
                    asks: Vec::new(),
                    last_updated: now,
                });
            let level = OrderbookLevel { price: a.price.clone(), size: size.clone() };
            match side.to_uppercase().as_str() {
                "BUY"  => apply_deltas(&mut entry.bids, std::slice::from_ref(&level), true),
                "SELL" => apply_deltas(&mut entry.asks, std::slice::from_ref(&level), false),
                _ => continue,
            }
            entry.last_updated = now;
        }
    }

    /// Get the full book for an asset.
    pub fn get_book(&self, asset_id: &str) -> Option<(&[OrderbookLevel], &[OrderbookLevel])> {
        self.books.get(asset_id).map(|e| (e.bids.as_slice(), e.asks.as_slice()))
    }

    /// Best bid (highest price).
    pub fn best_bid(&self, asset_id: &str) -> Option<&OrderbookLevel> {
        self.books.get(asset_id).and_then(|e| e.bids.first())
    }

    /// Best ask (lowest price).
    pub fn best_ask(&self, asset_id: &str) -> Option<&OrderbookLevel> {
        self.books.get(asset_id).and_then(|e| e.asks.first())
    }

    /// Spread between best ask and best bid.
    pub fn spread(&self, asset_id: &str) -> Option<f64> {
        let bid = self.best_bid(asset_id)?;
        let ask = self.best_ask(asset_id)?;
        let bid_price: f64 = bid.price.parse().ok()?;
        let ask_price: f64 = ask.price.parse().ok()?;
        Some(ask_price - bid_price)
    }

    /// Midpoint between best bid and best ask.
    pub fn midpoint(&self, asset_id: &str) -> Option<f64> {
        let bid = self.best_bid(asset_id)?;
        let ask = self.best_ask(asset_id)?;
        let bid_price: f64 = bid.price.parse().ok()?;
        let ask_price: f64 = ask.price.parse().ok()?;
        Some((bid_price + ask_price) / 2.0)
    }

    /// Number of tracked assets.
    pub fn len(&self) -> usize {
        self.books.len()
    }

    /// Whether any assets are tracked.
    pub fn is_empty(&self) -> bool {
        self.books.is_empty()
    }

    /// Clear all books.
    pub fn clear(&mut self) {
        self.books.clear();
    }

    /// All currently tracked asset IDs.
    pub fn tracked_tokens(&self) -> Vec<String> {
        self.books.keys().cloned().collect()
    }

    /// When the local copy of this asset was last touched (snapshot or delta).
    pub fn last_change(&self, asset_id: &str) -> Option<Instant> {
        self.books.get(asset_id).map(|e| e.last_updated)
    }

    /// All tracked assets whose last update is older than `threshold`.
    pub fn inactive_since(&self, threshold: Duration) -> Vec<String> {
        let now = Instant::now();
        self.books
            .iter()
            .filter(|(_, e)| now.duration_since(e.last_updated) >= threshold)
            .map(|(k, _)| k.clone())
            .collect()
    }

    /// Midpoints for the requested assets (missing tokens are omitted).
    pub fn midpoints(&self, asset_ids: &[String]) -> HashMap<String, f64> {
        asset_ids
            .iter()
            .filter_map(|id| self.midpoint(id).map(|m| (id.clone(), m)))
            .collect()
    }

    /// Spreads for the requested assets (missing tokens are omitted).
    pub fn spreads(&self, asset_ids: &[String]) -> HashMap<String, f64> {
        asset_ids
            .iter()
            .filter_map(|id| self.spread(id).map(|s| (id.clone(), s)))
            .collect()
    }

    /// Best bids for the requested assets (missing tokens are omitted).
    pub fn best_bids(&self, asset_ids: &[String]) -> HashMap<String, OrderbookLevel> {
        asset_ids
            .iter()
            .filter_map(|id| self.best_bid(id).cloned().map(|l| (id.clone(), l)))
            .collect()
    }

    /// Best asks for the requested assets (missing tokens are omitted).
    pub fn best_asks(&self, asset_ids: &[String]) -> HashMap<String, OrderbookLevel> {
        asset_ids
            .iter()
            .filter_map(|id| self.best_ask(id).cloned().map(|l| (id.clone(), l)))
            .collect()
    }

    /// Full books for the requested assets (missing tokens are omitted).
    pub fn books(&self, asset_ids: &[String]) -> HashMap<String, (Vec<OrderbookLevel>, Vec<OrderbookLevel>)> {
        asset_ids
            .iter()
            .filter_map(|id| {
                self.books
                    .get(id)
                    .map(|e| (id.clone(), (e.bids.clone(), e.asks.clone())))
            })
            .collect()
    }

    /// Midpoints for every tracked asset.
    pub fn midpoints_all(&self) -> HashMap<String, f64> {
        self.books
            .iter()
            .filter_map(|(id, _)| self.midpoint(id).map(|m| (id.clone(), m)))
            .collect()
    }

    /// Spreads for every tracked asset.
    pub fn spreads_all(&self) -> HashMap<String, f64> {
        self.books
            .iter()
            .filter_map(|(id, _)| self.spread(id).map(|s| (id.clone(), s)))
            .collect()
    }

    /// Best bids for every tracked asset.
    pub fn best_bids_all(&self) -> HashMap<String, OrderbookLevel> {
        self.books
            .iter()
            .filter_map(|(id, e)| e.bids.first().cloned().map(|l| (id.clone(), l)))
            .collect()
    }

    /// Best asks for every tracked asset.
    pub fn best_asks_all(&self) -> HashMap<String, OrderbookLevel> {
        self.books
            .iter()
            .filter_map(|(id, e)| e.asks.first().cloned().map(|l| (id.clone(), l)))
            .collect()
    }

    /// Full books for every tracked asset.
    pub fn books_all(&self) -> HashMap<String, (Vec<OrderbookLevel>, Vec<OrderbookLevel>)> {
        self.books
            .iter()
            .map(|(id, e)| (id.clone(), (e.bids.clone(), e.asks.clone())))
            .collect()
    }
}

impl Default for LocalOrderbook {
    fn default() -> Self {
        Self::new()
    }
}

fn apply_deltas(levels: &mut Vec<OrderbookLevel>, deltas: &[OrderbookLevel], descending: bool) {
    for delta in deltas {
        // Remove existing level at this price
        levels.retain(|l| l.price != delta.price);
        // Add if size is non-zero
        if delta.size != "0" && delta.size != "0.00" {
            levels.push(delta.clone());
        }
    }
    // Re-sort
    levels.sort_by(|a, b| {
        let pa: f64 = a.price.parse().unwrap_or(0.0);
        let pb: f64 = b.price.parse().unwrap_or(0.0);
        if descending {
            pb.partial_cmp(&pa).unwrap_or(std::cmp::Ordering::Equal)
        } else {
            pa.partial_cmp(&pb).unwrap_or(std::cmp::Ordering::Equal)
        }
    });
}

#[cfg(test)]
mod tests {
    use super::*;

    fn lvl(price: &str, size: &str) -> OrderbookLevel {
        OrderbookLevel {
            price: price.into(),
            size: size.into(),
        }
    }

    fn snap(asset: &str, bids: Vec<OrderbookLevel>, asks: Vec<OrderbookLevel>) -> BookSnapshot {
        BookSnapshot {
            asset_id: asset.into(),
            market: "m".into(),
            event_title: None,
            question: None,
            outcome: None,
            slug: None,
            bids,
            asks,
        }
    }

    fn upd(asset: &str, bids: Vec<OrderbookLevel>, asks: Vec<OrderbookLevel>) -> BookUpdate {
        BookUpdate {
            asset_id: asset.into(),
            market: "m".into(),
            event_title: None,
            question: None,
            outcome: None,
            slug: None,
            bids,
            asks,
        }
    }

    #[test]
    fn snapshot_stamps_last_change() {
        let mut book = LocalOrderbook::new();
        assert!(book.last_change("a").is_none());
        book.apply_snapshot(&snap("a", vec![lvl("0.5", "10")], vec![lvl("0.6", "10")]));
        assert!(book.last_change("a").is_some());
    }

    #[test]
    fn update_bumps_last_change() {
        let mut book = LocalOrderbook::new();
        book.apply_snapshot(&snap("a", vec![lvl("0.5", "10")], vec![lvl("0.6", "10")]));
        let first = book.last_change("a").unwrap();
        std::thread::sleep(Duration::from_millis(15));
        book.apply_update(&upd("a", vec![lvl("0.51", "5")], vec![]));
        let second = book.last_change("a").unwrap();
        assert!(second > first, "expected {:?} > {:?}", second, first);
    }

    #[test]
    fn inactive_since_filters_correctly() {
        let mut book = LocalOrderbook::new();
        book.apply_snapshot(&snap("fresh", vec![lvl("0.5", "10")], vec![lvl("0.6", "10")]));
        std::thread::sleep(Duration::from_millis(40));
        book.apply_snapshot(&snap("also_fresh", vec![lvl("0.4", "10")], vec![lvl("0.5", "10")]));

        let stale = book.inactive_since(Duration::from_millis(20));
        assert_eq!(stale, vec!["fresh".to_string()]);

        let none_stale = book.inactive_since(Duration::from_secs(60));
        assert!(none_stale.is_empty());
    }

    #[test]
    fn batch_methods_skip_missing_tokens() {
        let mut book = LocalOrderbook::new();
        book.apply_snapshot(&snap("a", vec![lvl("0.5", "10")], vec![lvl("0.6", "10")]));
        book.apply_snapshot(&snap("b", vec![lvl("0.4", "10")], vec![lvl("0.5", "10")]));

        let ids = vec!["a".to_string(), "b".to_string(), "missing".to_string()];
        let mids = book.midpoints(&ids);
        assert_eq!(mids.len(), 2);
        assert!(mids.contains_key("a"));
        assert!(mids.contains_key("b"));
        assert!(!mids.contains_key("missing"));

        let books = book.books(&ids);
        assert_eq!(books.len(), 2);
        assert!(books.contains_key("a"));
    }

    #[test]
    fn all_variants_cover_every_tracked_token() {
        let mut book = LocalOrderbook::new();
        book.apply_snapshot(&snap("a", vec![lvl("0.5", "10")], vec![lvl("0.6", "10")]));
        book.apply_snapshot(&snap("b", vec![lvl("0.4", "10")], vec![lvl("0.5", "10")]));

        assert_eq!(book.tracked_tokens().len(), 2);
        assert_eq!(book.midpoints_all().len(), 2);
        assert_eq!(book.spreads_all().len(), 2);
        assert_eq!(book.best_bids_all().len(), 2);
        assert_eq!(book.best_asks_all().len(), 2);
        assert_eq!(book.books_all().len(), 2);

        let mids = book.midpoints_all();
        assert!((mids["a"] - 0.55).abs() < 1e-9);
        assert!((mids["b"] - 0.45).abs() < 1e-9);
    }
}