use std::collections::HashMap;

use uforia_ms_types::{Asset, Venue};

use crate::L2OrderBook;

/// Aggregated order book across multiple venues for cross-venue analytics.
#[derive(Debug, Clone)]
pub struct CrossVenueBook {
    books: HashMap<(Asset, Venue), L2OrderBook>,
}

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

    /// Insert or replace the book for a given (asset, venue) pair.
    pub fn update_book(&mut self, asset: Asset, venue: Venue, book: &L2OrderBook) {
        self.books.insert((asset, venue), book.clone());
    }

    /// Cross-venue spread: max bid across all venues minus min ask across all venues
    /// for the given asset.
    ///
    /// A negative value indicates an arbitrage opportunity (crossed books).
    pub fn cross_venue_spread(&self, asset: Asset) -> Option<f64> {
        let mut max_bid: Option<f64> = None;
        let mut min_ask: Option<f64> = None;

        for ((a, _), book) in &self.books {
            if *a != asset {
                continue;
            }
            if let Some((bp, _)) = book.best_bid() {
                max_bid = Some(match max_bid {
                    Some(cur) => cur.max(bp),
                    None => bp,
                });
            }
            if let Some((ap, _)) = book.best_ask() {
                min_ask = Some(match min_ask {
                    Some(cur) => cur.min(ap),
                    None => ap,
                });
            }
        }

        match (max_bid, min_ask) {
            (Some(bid), Some(ask)) => Some(ask - bid),
            _ => None,
        }
    }

    /// Cross-venue depth imbalance for a given asset.
    ///
    /// Sums the top-of-book bid and ask sizes across all venues and returns
    /// (total_bid - total_ask) / (total_bid + total_ask).
    /// Returns `None` if no books are available for the asset.
    pub fn cross_venue_depth_imbalance(&self, asset: Asset) -> Option<f64> {
        let mut total_bid = 0.0;
        let mut total_ask = 0.0;
        let mut found = false;

        for ((a, _), book) in &self.books {
            if *a != asset {
                continue;
            }
            if let Some((_, bs)) = book.best_bid() {
                total_bid += bs;
                found = true;
            }
            if let Some((_, as_)) = book.best_ask() {
                total_ask += as_;
                found = true;
            }
        }

        if !found {
            return None;
        }

        let total = total_bid + total_ask;
        if total == 0.0 {
            return Some(0.0);
        }
        Some((total_bid - total_ask) / total)
    }
}

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

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

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

    fn make_book(bids: Vec<(f64, f64)>, asks: Vec<(f64, f64)>) -> L2OrderBook {
        let mut book = L2OrderBook::new();
        let snap = L2Snapshot {
            timestamp: Utc::now(),
            asset: Asset::Btc,
            venue: Venue::BinanceSpot, // venue in snapshot doesn't matter for the book itself
            bids: bids.into_iter().map(|(p, s)| make_level(p, s)).collect(),
            asks: asks.into_iter().map(|(p, s)| make_level(p, s)).collect(),
            last_update_id: None,
        };
        book.apply_snapshot(&snap);
        book
    }

    #[test]
    fn test_cross_venue_spread_normal() {
        let mut cv = CrossVenueBook::new();

        // Venue A: bid=50000, ask=50010
        let book_a = make_book(vec![(50000.0, 1.0)], vec![(50010.0, 1.0)]);
        cv.update_book(Asset::Btc, Venue::BinanceSpot, &book_a);

        // Venue B: bid=49995, ask=50015
        let book_b = make_book(vec![(49995.0, 1.0)], vec![(50015.0, 1.0)]);
        cv.update_book(Asset::Btc, Venue::BybitSpot, &book_b);

        // max bid = 50000 (A), min ask = 50010 (A), cross-venue spread = 10
        let spread = cv.cross_venue_spread(Asset::Btc).unwrap();
        assert!((spread - 10.0).abs() < 1e-10);
    }

    #[test]
    fn test_cross_venue_spread_crossed_arbitrage() {
        let mut cv = CrossVenueBook::new();

        // Venue A: bid=50020, ask=50030 (high market)
        let book_a = make_book(vec![(50020.0, 2.0)], vec![(50030.0, 1.0)]);
        cv.update_book(Asset::Btc, Venue::BinanceSpot, &book_a);

        // Venue B: bid=49990, ask=50010 (low market)
        let book_b = make_book(vec![(49990.0, 1.0)], vec![(50010.0, 2.0)]);
        cv.update_book(Asset::Btc, Venue::BybitSpot, &book_b);

        // max bid = 50020 (A), min ask = 50010 (B)
        // cross-venue spread = 50010 - 50020 = -10 (crossed = arb opportunity)
        let spread = cv.cross_venue_spread(Asset::Btc).unwrap();
        assert!((spread - (-10.0)).abs() < 1e-10);
    }

    #[test]
    fn test_cross_venue_depth_imbalance() {
        let mut cv = CrossVenueBook::new();

        // Venue A: bid_size=3, ask_size=1
        let book_a = make_book(vec![(50000.0, 3.0)], vec![(50010.0, 1.0)]);
        cv.update_book(Asset::Btc, Venue::BinanceSpot, &book_a);

        // Venue B: bid_size=2, ask_size=4
        let book_b = make_book(vec![(49995.0, 2.0)], vec![(50015.0, 4.0)]);
        cv.update_book(Asset::Btc, Venue::BybitSpot, &book_b);

        // total_bid = 3+2 = 5, total_ask = 1+4 = 5
        // imbalance = (5 - 5) / (5 + 5) = 0
        let imb = cv.cross_venue_depth_imbalance(Asset::Btc).unwrap();
        assert!((imb - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_cross_venue_depth_imbalance_skewed() {
        let mut cv = CrossVenueBook::new();

        // Venue A: bid_size=5, ask_size=1
        let book_a = make_book(vec![(50000.0, 5.0)], vec![(50010.0, 1.0)]);
        cv.update_book(Asset::Btc, Venue::BinanceSpot, &book_a);

        // Venue B: bid_size=5, ask_size=1
        let book_b = make_book(vec![(49995.0, 5.0)], vec![(50015.0, 1.0)]);
        cv.update_book(Asset::Btc, Venue::BybitSpot, &book_b);

        // total_bid = 10, total_ask = 2
        // imbalance = (10 - 2) / (10 + 2) = 8/12 = 2/3
        let imb = cv.cross_venue_depth_imbalance(Asset::Btc).unwrap();
        assert!((imb - (2.0 / 3.0)).abs() < 1e-10);
    }

    #[test]
    fn test_no_data_for_asset() {
        let cv = CrossVenueBook::new();
        assert!(cv.cross_venue_spread(Asset::Btc).is_none());
        assert!(cv.cross_venue_depth_imbalance(Asset::Btc).is_none());
    }

    #[test]
    fn test_asset_isolation() {
        let mut cv = CrossVenueBook::new();

        let btc_book = make_book(vec![(50000.0, 1.0)], vec![(50010.0, 1.0)]);
        cv.update_book(Asset::Btc, Venue::BinanceSpot, &btc_book);

        let eth_book = make_book(vec![(3000.0, 10.0)], vec![(3001.0, 10.0)]);
        cv.update_book(Asset::Eth, Venue::BinanceSpot, &eth_book);

        // BTC spread should not be affected by ETH data
        let btc_spread = cv.cross_venue_spread(Asset::Btc).unwrap();
        assert!((btc_spread - 10.0).abs() < 1e-10);

        let eth_spread = cv.cross_venue_spread(Asset::Eth).unwrap();
        assert!((eth_spread - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_update_replaces_book() {
        let mut cv = CrossVenueBook::new();

        let book_v1 = make_book(vec![(50000.0, 1.0)], vec![(50010.0, 1.0)]);
        cv.update_book(Asset::Btc, Venue::BinanceSpot, &book_v1);

        // Replace with tighter spread
        let book_v2 = make_book(vec![(50005.0, 2.0)], vec![(50006.0, 2.0)]);
        cv.update_book(Asset::Btc, Venue::BinanceSpot, &book_v2);

        let spread = cv.cross_venue_spread(Asset::Btc).unwrap();
        assert!((spread - 1.0).abs() < 1e-10);
    }
}
