//! L2 orderbook snapshot with derived metrics.
//!
//! All metric computations are `#[inline]` for hot-path use.
//! Levels stored as Vec<(f64, f64)> — cache-friendly, zero-copy to Python.

use pyo3::prelude::*;

/// Full L2 orderbook snapshot for a single instrument.
///
/// Bids sorted descending by price, asks sorted ascending by price.
#[pyclass]
#[derive(Debug, Clone)]
pub struct OrderbookSnapshot {
    /// Bid levels: Vec<(price, size)>, sorted descending by price.
    pub(crate) bids: Vec<(f64, f64)>,
    /// Ask levels: Vec<(price, size)>, sorted ascending by price.
    pub(crate) asks: Vec<(f64, f64)>,
    #[pyo3(get)]
    pub timestamp: f64,
    #[pyo3(get)]
    pub source: String,
}

impl Default for OrderbookSnapshot {
    fn default() -> Self {
        Self {
            bids: Vec::new(),
            asks: Vec::new(),
            timestamp: 0.0,
            source: String::new(),
        }
    }
}

impl OrderbookSnapshot {
    /// Create from raw level data (Rust-side constructor).
    pub fn from_levels(
        bids: Vec<(f64, f64)>,
        asks: Vec<(f64, f64)>,
        timestamp: f64,
        source: String,
    ) -> Self {
        Self {
            bids,
            asks,
            timestamp,
            source,
        }
    }
}

#[pymethods]
impl OrderbookSnapshot {
    #[new]
    #[pyo3(signature = (timestamp=0.0, source=""))]
    fn new(timestamp: f64, source: &str) -> Self {
        Self {
            bids: Vec::new(),
            asks: Vec::new(),
            timestamp,
            source: source.to_string(),
        }
    }

    /// All bid levels as list of (price, size) tuples, descending by price.
    fn bids(&self) -> Vec<(f64, f64)> {
        self.bids.clone()
    }

    /// All ask levels as list of (price, size) tuples, ascending by price.
    fn asks(&self) -> Vec<(f64, f64)> {
        self.asks.clone()
    }

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

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

    /// Best bid (price, size) or None.
    fn best_bid(&self) -> Option<(f64, f64)> {
        self.bids.first().copied()
    }

    /// Best ask (price, size) or None.
    fn best_ask(&self) -> Option<(f64, f64)> {
        self.asks.first().copied()
    }

    /// Spread: best_ask - best_bid. Returns 0.0 if no levels.
    #[inline]
    fn spread(&self) -> f64 {
        match (self.bids.first(), self.asks.first()) {
            (Some(&(bid, _)), Some(&(ask, _))) => ask - bid,
            _ => 0.0,
        }
    }

    /// Mid price: (best_bid + best_ask) / 2. Returns 0.0 if no levels.
    #[inline]
    fn mid(&self) -> f64 {
        match (self.bids.first(), self.asks.first()) {
            (Some(&(bid, _)), Some(&(ask, _))) => (bid + ask) / 2.0,
            _ => 0.0,
        }
    }

    /// Volume-weighted mid price using top N levels.
    #[inline]
    #[pyo3(signature = (levels=5))]
    fn vwap_mid(&self, levels: usize) -> f64 {
        let bid_vwap = vwap_side(&self.bids, levels);
        let ask_vwap = vwap_side(&self.asks, levels);
        match (bid_vwap, ask_vwap) {
            (Some(b), Some(a)) => (b + a) / 2.0,
            (Some(b), None) => b,
            (None, Some(a)) => a,
            (None, None) => 0.0,
        }
    }

    /// Order book imbalance: bid_size / (bid_size + ask_size) for top N levels.
    /// Returns 0.5 (balanced) if no data. >0.5 = buy pressure, <0.5 = sell pressure.
    #[inline]
    #[pyo3(signature = (levels=5))]
    fn imbalance(&self, levels: usize) -> f64 {
        let bid_total = total_size(&self.bids, levels);
        let ask_total = total_size(&self.asks, levels);
        let sum = bid_total + ask_total;
        if sum <= 0.0 {
            0.5
        } else {
            bid_total / sum
        }
    }

    /// Total bid depth (size) across top N levels.
    #[inline]
    #[pyo3(signature = (levels=None))]
    fn bid_depth(&self, levels: Option<usize>) -> f64 {
        total_size(&self.bids, levels.unwrap_or(self.bids.len()))
    }

    /// Total ask depth (size) across top N levels.
    #[inline]
    #[pyo3(signature = (levels=None))]
    fn ask_depth(&self, levels: Option<usize>) -> f64 {
        total_size(&self.asks, levels.unwrap_or(self.asks.len()))
    }

    /// Cumulative depth at a given price deviation from mid.
    /// Returns (bid_depth_within, ask_depth_within).
    #[inline]
    fn depth_within(&self, price_range: f64) -> (f64, f64) {
        let mid = self.mid();
        if mid <= 0.0 {
            return (0.0, 0.0);
        }
        let bid_depth: f64 = self
            .bids
            .iter()
            .take_while(|&&(p, _)| mid - p <= price_range)
            .map(|&(_, s)| s)
            .sum();
        let ask_depth: f64 = self
            .asks
            .iter()
            .take_while(|&&(p, _)| p - mid <= price_range)
            .map(|&(_, s)| s)
            .sum();
        (bid_depth, ask_depth)
    }

    fn __repr__(&self) -> String {
        format!(
            "OrderbookSnapshot(bids={} asks={} spread={:.4} mid={:.4})",
            self.bids.len(),
            self.asks.len(),
            self.spread(),
            self.mid(),
        )
    }
}

/// VWAP for one side of the book.
#[inline]
fn vwap_side(levels: &[(f64, f64)], n: usize) -> Option<f64> {
    let mut pv_sum = 0.0;
    let mut v_sum = 0.0;
    for &(price, size) in levels.iter().take(n) {
        pv_sum += price * size;
        v_sum += size;
    }
    if v_sum > 0.0 {
        Some(pv_sum / v_sum)
    } else {
        None
    }
}

/// Total size for top N levels.
#[inline]
fn total_size(levels: &[(f64, f64)], n: usize) -> f64 {
    levels.iter().take(n).map(|&(_, s)| s).sum()
}

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

    fn sample_book() -> OrderbookSnapshot {
        OrderbookSnapshot::from_levels(
            vec![(0.55, 100.0), (0.54, 200.0), (0.53, 50.0)],
            vec![(0.56, 80.0), (0.57, 150.0), (0.58, 300.0)],
            1000.0,
            "test".to_string(),
        )
    }

    #[test]
    fn best_bid_ask() {
        let book = sample_book();
        assert_eq!(book.best_bid(), Some((0.55, 100.0)));
        assert_eq!(book.best_ask(), Some((0.56, 80.0)));
    }

    #[test]
    fn spread() {
        let book = sample_book();
        assert!((book.spread() - 0.01).abs() < 1e-10);
    }

    #[test]
    fn mid() {
        let book = sample_book();
        assert!((book.mid() - 0.555).abs() < 1e-10);
    }

    #[test]
    fn imbalance_balanced() {
        let book = OrderbookSnapshot::from_levels(
            vec![(0.55, 100.0)],
            vec![(0.56, 100.0)],
            0.0,
            String::new(),
        );
        assert!((book.imbalance(1) - 0.5).abs() < 1e-10);
    }

    #[test]
    fn imbalance_bid_heavy() {
        let book = OrderbookSnapshot::from_levels(
            vec![(0.55, 300.0)],
            vec![(0.56, 100.0)],
            0.0,
            String::new(),
        );
        assert!(book.imbalance(1) > 0.5);
        assert!((book.imbalance(1) - 0.75).abs() < 1e-10);
    }

    #[test]
    fn bid_ask_depth() {
        let book = sample_book();
        assert!((book.bid_depth(Some(2)) - 300.0).abs() < 1e-10);
        assert!((book.ask_depth(Some(2)) - 230.0).abs() < 1e-10);
        assert!((book.bid_depth(None) - 350.0).abs() < 1e-10);
    }

    #[test]
    fn vwap_mid() {
        let book = sample_book();
        let vwap = book.vwap_mid(2);
        // bid VWAP(2) = (0.55*100 + 0.54*200) / 300 = 163/300 = 0.54333...
        // ask VWAP(2) = (0.56*80 + 0.57*150) / 230 = 130.3/230 = 0.56652...
        // vwap_mid = (0.54333 + 0.56652) / 2 ≈ 0.55493
        assert!(vwap > 0.54 && vwap < 0.57);
    }

    #[test]
    fn empty_book() {
        let book = OrderbookSnapshot::default();
        assert_eq!(book.spread(), 0.0);
        assert_eq!(book.mid(), 0.0);
        assert_eq!(book.imbalance(5), 0.5);
        assert_eq!(book.best_bid(), None);
        assert_eq!(book.best_ask(), None);
    }

    #[test]
    fn depth_within() {
        let book = sample_book();
        // mid = 0.555
        // bids within 0.02 of mid: 0.55 (100), 0.54 (200) — 0.555-0.55=0.005, 0.555-0.54=0.015
        // asks within 0.02 of mid: 0.56 (80), 0.57 (150) — 0.56-0.555=0.005, 0.57-0.555=0.015
        let (bd, ad) = book.depth_within(0.02);
        assert!((bd - 300.0).abs() < 1e-10);
        assert!((ad - 230.0).abs() < 1e-10);
    }
}
