use std::collections::BTreeMap;

use ordered_float::OrderedFloat;
use tracing::trace;

use uforia_ms_types::{BookLevel, L2Snapshot, L2Update};

/// An L2 (price-level) order book backed by sorted trees.
///
/// Bids are stored in a `BTreeMap` and iterated in reverse (highest first).
/// Asks are stored in natural order (lowest first).
#[derive(Debug, Clone)]
pub struct L2OrderBook {
    /// bid price -> size (highest price = best bid)
    bids: BTreeMap<OrderedFloat<f64>, f64>,
    /// ask price -> size (lowest price = best ask)
    asks: BTreeMap<OrderedFloat<f64>, f64>,
}

impl L2OrderBook {
    pub fn new() -> Self {
        Self {
            bids: BTreeMap::new(),
            asks: BTreeMap::new(),
        }
    }

    /// Replace the entire book with the given snapshot.
    pub fn apply_snapshot(&mut self, snap: &L2Snapshot) {
        self.bids.clear();
        self.asks.clear();
        for level in &snap.bids {
            if level.size > 0.0 {
                self.bids.insert(level.price, level.size);
            }
        }
        for level in &snap.asks {
            if level.size > 0.0 {
                self.asks.insert(level.price, level.size);
            }
        }
        trace!(
            bids = self.bids.len(),
            asks = self.asks.len(),
            "applied L2 snapshot"
        );
    }

    /// Apply an incremental update. Levels with size == 0 are removed.
    pub fn apply_update(&mut self, update: &L2Update) {
        for level in &update.bids {
            apply_level(&mut self.bids, level);
        }
        for level in &update.asks {
            apply_level(&mut self.asks, level);
        }
        trace!(
            bid_updates = update.bids.len(),
            ask_updates = update.asks.len(),
            "applied L2 update"
        );
    }

    /// Best bid: highest bid price and its size.
    pub fn best_bid(&self) -> Option<(f64, f64)> {
        self.bids
            .iter()
            .next_back()
            .map(|(p, s)| (p.into_inner(), *s))
    }

    /// Best ask: lowest ask price and its size.
    pub fn best_ask(&self) -> Option<(f64, f64)> {
        self.asks.iter().next().map(|(p, s)| (p.into_inner(), *s))
    }

    /// Midpoint price: (best_bid + best_ask) / 2.
    pub fn mid(&self) -> Option<f64> {
        match (self.best_bid(), self.best_ask()) {
            (Some((bid, _)), Some((ask, _))) => Some((bid + ask) / 2.0),
            _ => None,
        }
    }

    /// Size-weighted microprice from BBO.
    ///
    /// microprice = (bid_price * ask_size + ask_price * bid_size) / (bid_size + ask_size)
    pub fn microprice(&self) -> Option<f64> {
        match (self.best_bid(), self.best_ask()) {
            (Some((bp, bs)), Some((ap, as_))) => {
                let total = bs + as_;
                if total == 0.0 {
                    return self.mid();
                }
                Some((bp * as_ + ap * bs) / total)
            }
            _ => None,
        }
    }

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

    /// Relative spread: spread / mid.
    pub fn spread_rel(&self) -> Option<f64> {
        match (self.spread(), self.mid()) {
            (Some(s), Some(m)) if m != 0.0 => Some(s / m),
            _ => None,
        }
    }

    /// Sum of the top `n` bid level sizes.
    pub fn depth_bid(&self, n: usize) -> f64 {
        self.bids.iter().rev().take(n).map(|(_, s)| s).sum()
    }

    /// Sum of the top `n` ask level sizes.
    pub fn depth_ask(&self, n: usize) -> f64 {
        self.asks.iter().take(n).map(|(_, s)| s).sum()
    }

    /// Order-book imbalance over top `n` levels.
    ///
    /// Returns (bid_depth - ask_depth) / (bid_depth + ask_depth).
    /// Returns 0.0 if both sides are empty.
    pub fn imbalance(&self, n: usize) -> f64 {
        let bd = self.depth_bid(n);
        let ad = self.depth_ask(n);
        let total = bd + ad;
        if total == 0.0 {
            return 0.0;
        }
        (bd - ad) / total
    }

    /// Number of bid levels.
    pub fn bid_count(&self) -> usize {
        self.bids.len()
    }

    /// Number of ask levels.
    pub fn ask_count(&self) -> usize {
        self.asks.len()
    }
}

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

fn apply_level(side: &mut BTreeMap<OrderedFloat<f64>, f64>, level: &BookLevel) {
    if level.size == 0.0 {
        side.remove(&level.price);
    } else {
        side.insert(level.price, level.size);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;
    use ordered_float::OrderedFloat;
    use uforia_ms_types::{Asset, Venue};

    fn make_level(price: f64, size: f64) -> BookLevel {
        BookLevel {
            price: OrderedFloat(price),
            size,
        }
    }

    fn sample_snapshot() -> L2Snapshot {
        L2Snapshot {
            timestamp: Utc::now(),
            asset: Asset::Btc,
            venue: Venue::BinanceSpot,
            bids: vec![
                make_level(50000.0, 1.0),
                make_level(49990.0, 2.0),
                make_level(49980.0, 3.0),
            ],
            asks: vec![
                make_level(50010.0, 1.5),
                make_level(50020.0, 2.5),
                make_level(50030.0, 3.5),
            ],
            last_update_id: Some(100),
        }
    }

    #[test]
    fn test_apply_snapshot_bbo() {
        let mut book = L2OrderBook::new();
        book.apply_snapshot(&sample_snapshot());

        let (bp, bs) = book.best_bid().unwrap();
        assert!((bp - 50000.0).abs() < 1e-10);
        assert!((bs - 1.0).abs() < 1e-10);

        let (ap, as_) = book.best_ask().unwrap();
        assert!((ap - 50010.0).abs() < 1e-10);
        assert!((as_ - 1.5).abs() < 1e-10);

        assert_eq!(book.bid_count(), 3);
        assert_eq!(book.ask_count(), 3);
    }

    #[test]
    fn test_mid_and_spread() {
        let mut book = L2OrderBook::new();
        book.apply_snapshot(&sample_snapshot());

        let mid = book.mid().unwrap();
        assert!((mid - 50005.0).abs() < 1e-10);

        let spread = book.spread().unwrap();
        assert!((spread - 10.0).abs() < 1e-10);

        let spread_rel = book.spread_rel().unwrap();
        assert!((spread_rel - 10.0 / 50005.0).abs() < 1e-10);
    }

    #[test]
    fn test_apply_update_modifies_level() {
        let mut book = L2OrderBook::new();
        book.apply_snapshot(&sample_snapshot());

        let update = L2Update {
            timestamp: Utc::now(),
            asset: Asset::Btc,
            venue: Venue::BinanceSpot,
            bids: vec![make_level(50000.0, 5.0)],
            asks: vec![make_level(50010.0, 0.5)],
            first_update_id: Some(101),
            last_update_id: Some(101),
        };
        book.apply_update(&update);

        let (_, bs) = book.best_bid().unwrap();
        assert!((bs - 5.0).abs() < 1e-10);

        let (_, as_) = book.best_ask().unwrap();
        assert!((as_ - 0.5).abs() < 1e-10);
    }

    #[test]
    fn test_apply_update_delete_level() {
        let mut book = L2OrderBook::new();
        book.apply_snapshot(&sample_snapshot());
        assert_eq!(book.bid_count(), 3);

        // Delete the best bid by setting size to 0
        let update = L2Update {
            timestamp: Utc::now(),
            asset: Asset::Btc,
            venue: Venue::BinanceSpot,
            bids: vec![make_level(50000.0, 0.0)],
            asks: vec![],
            first_update_id: Some(102),
            last_update_id: Some(102),
        };
        book.apply_update(&update);

        assert_eq!(book.bid_count(), 2);
        let (bp, _) = book.best_bid().unwrap();
        assert!((bp - 49990.0).abs() < 1e-10);
    }

    #[test]
    fn test_microprice_symmetric() {
        let mut book = L2OrderBook::new();
        let snap = L2Snapshot {
            timestamp: Utc::now(),
            asset: Asset::Btc,
            venue: Venue::BinanceSpot,
            bids: vec![make_level(100.0, 1.0)],
            asks: vec![make_level(102.0, 1.0)],
            last_update_id: None,
        };
        book.apply_snapshot(&snap);

        // Equal sizes -> microprice == mid
        let mp = book.microprice().unwrap();
        let mid = book.mid().unwrap();
        assert!((mp - mid).abs() < 1e-10);
    }

    #[test]
    fn test_microprice_asymmetric() {
        let mut book = L2OrderBook::new();
        let snap = L2Snapshot {
            timestamp: Utc::now(),
            asset: Asset::Btc,
            venue: Venue::BinanceSpot,
            bids: vec![make_level(100.0, 3.0)],
            asks: vec![make_level(102.0, 1.0)],
            last_update_id: None,
        };
        book.apply_snapshot(&snap);

        // microprice = (100*1 + 102*3) / (3+1) = 406/4 = 101.5
        let mp = book.microprice().unwrap();
        assert!((mp - 101.5).abs() < 1e-10);
    }

    #[test]
    fn test_depth_and_imbalance() {
        let mut book = L2OrderBook::new();
        book.apply_snapshot(&sample_snapshot());

        // Top 2 bids: 1.0 + 2.0 = 3.0
        assert!((book.depth_bid(2) - 3.0).abs() < 1e-10);
        // Top 2 asks: 1.5 + 2.5 = 4.0
        assert!((book.depth_ask(2) - 4.0).abs() < 1e-10);

        // imbalance(2) = (3.0 - 4.0) / (3.0 + 4.0) = -1/7
        let imb = book.imbalance(2);
        assert!((imb - (-1.0 / 7.0)).abs() < 1e-10);
    }

    #[test]
    fn test_depth_all_levels() {
        let mut book = L2OrderBook::new();
        book.apply_snapshot(&sample_snapshot());

        // All 3 bids: 1.0 + 2.0 + 3.0 = 6.0
        assert!((book.depth_bid(10) - 6.0).abs() < 1e-10);
        // All 3 asks: 1.5 + 2.5 + 3.5 = 7.5
        assert!((book.depth_ask(10) - 7.5).abs() < 1e-10);
    }

    #[test]
    fn test_empty_book() {
        let book = L2OrderBook::new();
        assert!(book.best_bid().is_none());
        assert!(book.best_ask().is_none());
        assert!(book.mid().is_none());
        assert!(book.microprice().is_none());
        assert!(book.spread().is_none());
        assert!(book.spread_rel().is_none());
        assert!((book.depth_bid(5) - 0.0).abs() < 1e-10);
        assert!((book.imbalance(5) - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_snapshot_replaces_previous() {
        let mut book = L2OrderBook::new();
        book.apply_snapshot(&sample_snapshot());
        assert_eq!(book.bid_count(), 3);

        // Apply a different snapshot with fewer levels
        let snap2 = L2Snapshot {
            timestamp: Utc::now(),
            asset: Asset::Btc,
            venue: Venue::BinanceSpot,
            bids: vec![make_level(60000.0, 0.5)],
            asks: vec![make_level(60010.0, 0.5)],
            last_update_id: Some(200),
        };
        book.apply_snapshot(&snap2);

        assert_eq!(book.bid_count(), 1);
        assert_eq!(book.ask_count(), 1);
        let (bp, _) = book.best_bid().unwrap();
        assert!((bp - 60000.0).abs() < 1e-10);
    }
}

/// Property-based tests using proptest.
#[cfg(test)]
mod proptests {
    use super::*;
    use chrono::Utc;
    use ordered_float::OrderedFloat;
    use proptest::prelude::*;
    use uforia_ms_types::{Asset, Venue};

    fn arb_levels(max_count: usize) -> impl Strategy<Value = Vec<BookLevel>> {
        prop::collection::vec(
            (1.0f64..100000.0, 0.0f64..1000.0).prop_map(|(p, s)| BookLevel {
                price: OrderedFloat(p),
                size: s,
            }),
            0..max_count,
        )
    }

    proptest! {
        #[test]
        fn prop_no_negative_sizes(
            bids in arb_levels(20),
            asks in arb_levels(20),
        ) {
            let mut book = L2OrderBook::new();
            let snap = L2Snapshot {
                timestamp: Utc::now(),
                asset: Asset::Btc,
                venue: Venue::BinanceSpot,
                bids,
                asks,
                last_update_id: None,
            };
            book.apply_snapshot(&snap);

            // All sizes in the book must be non-negative
            for (_, size) in &book.bids {
                prop_assert!(*size >= 0.0, "negative bid size: {}", size);
            }
            for (_, size) in &book.asks {
                prop_assert!(*size >= 0.0, "negative ask size: {}", size);
            }

            // Depth must be non-negative
            prop_assert!(book.depth_bid(5) >= 0.0);
            prop_assert!(book.depth_ask(5) >= 0.0);

            // Imbalance in [-1, 1]
            let imb = book.imbalance(5);
            prop_assert!(imb >= -1.0 && imb <= 1.0, "imbalance out of range: {}", imb);
        }

        #[test]
        fn prop_spread_non_negative(
            bid_price in 1.0f64..50000.0,
            ask_offset in 0.01f64..1000.0,
            bid_size in 0.01f64..100.0,
            ask_size in 0.01f64..100.0,
        ) {
            let mut book = L2OrderBook::new();
            let ask_price = bid_price + ask_offset;
            let snap = L2Snapshot {
                timestamp: Utc::now(),
                asset: Asset::Btc,
                venue: Venue::BinanceSpot,
                bids: vec![BookLevel { price: OrderedFloat(bid_price), size: bid_size }],
                asks: vec![BookLevel { price: OrderedFloat(ask_price), size: ask_size }],
                last_update_id: None,
            };
            book.apply_snapshot(&snap);

            let spread = book.spread().unwrap();
            prop_assert!(spread >= 0.0, "negative spread: {}", spread);

            let mid = book.mid().unwrap();
            prop_assert!(mid >= bid_price);
            prop_assert!(mid <= ask_price);
        }
    }
}
