//! Cross-venue liquidity aggregation.
//!
//! Merges orderbooks from multiple venues, computes unified best bid/ask,
//! and routes orders to the best execution venue.

use pyo3::prelude::*;

/// A single price level in the aggregated book.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct AggregatedLevel {
    pub price: f64,
    pub size: f64,
    pub venue: String,
}

#[pymethods]
impl AggregatedLevel {
    #[new]
    fn new(price: f64, size: f64, venue: String) -> Self {
        Self { price, size, venue }
    }

    fn __repr__(&self) -> String {
        format!("AggregatedLevel(price={:.4}, size={:.2}, venue='{}')", self.price, self.size, self.venue)
    }
}

/// Aggregated orderbook across venues.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct AggregatedBook {
    pub bids: Vec<AggregatedLevel>,
    pub asks: Vec<AggregatedLevel>,
    pub best_bid: f64,
    pub best_ask: f64,
    pub mid_price: f64,
    pub total_bid_size: f64,
    pub total_ask_size: f64,
}

#[pymethods]
impl AggregatedBook {
    fn __repr__(&self) -> String {
        format!(
            "AggregatedBook(bid={:.4}, ask={:.4}, mid={:.4}, bid_depth={:.2}, ask_depth={:.2})",
            self.best_bid, self.best_ask, self.mid_price, self.total_bid_size, self.total_ask_size
        )
    }

    fn spread(&self) -> f64 {
        finite_or_zero(self.best_ask - self.best_bid)
    }
}

/// Smart routing decision.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct SmartRouteDecision {
    pub venue: String,
    pub price: f64,
    pub size: f64,
    pub expected_cost: f64,
    pub fills_across_venues: Vec<(String, f64, f64)>, // (venue, price, size)
}

#[pymethods]
impl SmartRouteDecision {
    fn __repr__(&self) -> String {
        format!("SmartRouteDecision(venue='{}', price={:.4}, cost={:.4})",
            self.venue, self.price, self.expected_cost)
    }
}

/// Venue mapping for cross-venue equivalence.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct VenueMapping {
    pub canonical_id: String,
    pub venue_market_ids: Vec<(String, String)>, // (venue_name, market_id)
}

#[pymethods]
impl VenueMapping {
    #[new]
    fn new(canonical_id: String, venue_market_ids: Vec<(String, String)>) -> Self {
        Self { canonical_id, venue_market_ids }
    }

    fn __repr__(&self) -> String {
        format!("VenueMapping(id='{}', venues={})", self.canonical_id, self.venue_market_ids.len())
    }
}

#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

/// Aggregate orderbooks from multiple venues into a unified book.
///
/// `venue_bids` and `venue_asks` are lists of (venue_name, price, size) tuples.
/// Returns sorted aggregated book (bids descending, asks ascending).
#[pyfunction]
pub fn aggregate_books(
    venue_bids: Vec<(String, f64, f64)>,
    venue_asks: Vec<(String, f64, f64)>,
) -> AggregatedBook {
    let mut bids: Vec<AggregatedLevel> = venue_bids
        .into_iter()
        .filter(|(_, p, s)| *p > 0.0 && *s > 0.0 && p.is_finite() && s.is_finite())
        .map(|(v, p, s)| AggregatedLevel { price: p, size: s, venue: v })
        .collect();

    let mut asks: Vec<AggregatedLevel> = venue_asks
        .into_iter()
        .filter(|(_, p, s)| *p > 0.0 && *s > 0.0 && p.is_finite() && s.is_finite())
        .map(|(v, p, s)| AggregatedLevel { price: p, size: s, venue: v })
        .collect();

    // Sort bids descending by price
    bids.sort_by(|a, b| b.price.partial_cmp(&a.price).unwrap_or(std::cmp::Ordering::Equal));
    // Sort asks ascending by price
    asks.sort_by(|a, b| a.price.partial_cmp(&b.price).unwrap_or(std::cmp::Ordering::Equal));

    let best_bid = bids.first().map(|l| l.price).unwrap_or(0.0);
    let best_ask = asks.first().map(|l| l.price).unwrap_or(0.0);
    let mid = if best_bid > 0.0 && best_ask > 0.0 {
        finite_or_zero((best_bid + best_ask) / 2.0)
    } else {
        0.0
    };
    let total_bid: f64 = bids.iter().map(|l| l.size).sum();
    let total_ask: f64 = asks.iter().map(|l| l.size).sum();

    AggregatedBook {
        bids,
        asks,
        best_bid,
        best_ask,
        mid_price: mid,
        total_bid_size: total_bid,
        total_ask_size: total_ask,
    }
}

/// Route an order to the best venue(s) based on aggregated book.
///
/// `side`: "buy" or "sell"
/// `size`: desired order size
/// Walks the aggregated book to find best execution.
#[pyfunction]
pub fn smart_route(
    book: &AggregatedBook,
    side: &str,
    size: f64,
) -> SmartRouteDecision {
    if size <= 0.0 {
        return SmartRouteDecision {
            venue: String::new(), price: 0.0, size: 0.0,
            expected_cost: 0.0, fills_across_venues: Vec::new(),
        };
    }

    let levels = if side == "buy" { &book.asks } else { &book.bids };
    let mut remaining = size;
    let mut total_cost = 0.0;
    let mut fills = Vec::new();

    for level in levels {
        if remaining <= 0.0 {
            break;
        }
        let fill_size = remaining.min(level.size);
        total_cost += fill_size * level.price;
        fills.push((level.venue.clone(), level.price, fill_size));
        remaining -= fill_size;
    }

    let filled_size = size - remaining;
    let avg_price = if filled_size > 0.0 {
        finite_or_zero(total_cost / filled_size)
    } else {
        0.0
    };

    let primary_venue = fills.first()
        .map(|(v, _, _)| v.clone())
        .unwrap_or_default();

    SmartRouteDecision {
        venue: primary_venue,
        price: avg_price,
        size: filled_size,
        expected_cost: finite_or_zero(total_cost),
        fills_across_venues: fills,
    }
}

/// Find the single best venue for an order.
#[pyfunction]
pub fn best_venue(
    book: &AggregatedBook,
    side: &str,
) -> (String, f64) {
    let levels = if side == "buy" { &book.asks } else { &book.bids };
    levels
        .first()
        .map(|l| (l.venue.clone(), l.price))
        .unwrap_or_default()
}

/// Compute the merged mid-price from aggregated book.
#[pyfunction]
pub fn merged_mid_price(
    venue_bids: Vec<(String, f64, f64)>,
    venue_asks: Vec<(String, f64, f64)>,
) -> f64 {
    let book = aggregate_books(venue_bids, venue_asks);
    book.mid_price
}

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<AggregatedLevel>()?;
    m.add_class::<AggregatedBook>()?;
    m.add_class::<SmartRouteDecision>()?;
    m.add_class::<VenueMapping>()?;
    m.add_function(wrap_pyfunction!(aggregate_books, m)?)?;
    m.add_function(wrap_pyfunction!(smart_route, m)?)?;
    m.add_function(wrap_pyfunction!(best_venue, m)?)?;
    m.add_function(wrap_pyfunction!(merged_mid_price, m)?)?;
    Ok(())
}

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

    #[test]
    fn test_aggregate_books_basic() {
        let bids = vec![
            ("poly".into(), 0.50, 100.0),
            ("kalshi".into(), 0.48, 200.0),
        ];
        let asks = vec![
            ("poly".into(), 0.52, 150.0),
            ("kalshi".into(), 0.55, 100.0),
        ];
        let book = aggregate_books(bids, asks);
        assert!((book.best_bid - 0.50).abs() < 1e-10);
        assert!((book.best_ask - 0.52).abs() < 1e-10);
        assert!((book.mid_price - 0.51).abs() < 1e-10);
    }

    #[test]
    fn test_aggregate_books_empty() {
        let book = aggregate_books(vec![], vec![]);
        assert_eq!(book.best_bid, 0.0);
        assert_eq!(book.best_ask, 0.0);
        assert_eq!(book.mid_price, 0.0);
    }

    #[test]
    fn test_smart_route_buy() {
        let book = aggregate_books(
            vec![("A".into(), 0.50, 100.0)],
            vec![
                ("A".into(), 0.52, 50.0),
                ("B".into(), 0.53, 100.0),
            ],
        );
        let decision = smart_route(&book, "buy", 80.0);
        assert_eq!(decision.fills_across_venues.len(), 2);
        assert!((decision.size - 80.0).abs() < 1e-10);
        assert!(decision.price > 0.52 && decision.price < 0.53);
    }

    #[test]
    fn test_smart_route_sell() {
        let book = aggregate_books(
            vec![
                ("A".into(), 0.50, 50.0),
                ("B".into(), 0.48, 100.0),
            ],
            vec![("A".into(), 0.52, 100.0)],
        );
        let decision = smart_route(&book, "sell", 30.0);
        assert_eq!(decision.venue, "A");
        assert!((decision.price - 0.50).abs() < 1e-10);
    }

    #[test]
    fn test_smart_route_zero_size() {
        let book = aggregate_books(vec![], vec![]);
        let decision = smart_route(&book, "buy", 0.0);
        assert!((decision.size).abs() < 1e-10);
    }

    #[test]
    fn test_best_venue_buy() {
        let book = aggregate_books(
            vec![],
            vec![
                ("poly".into(), 0.52, 100.0),
                ("kalshi".into(), 0.50, 50.0),
            ],
        );
        let (venue, price) = best_venue(&book, "buy");
        assert_eq!(venue, "kalshi");
        assert!((price - 0.50).abs() < 1e-10);
    }

    #[test]
    fn test_best_venue_sell() {
        let book = aggregate_books(
            vec![
                ("poly".into(), 0.50, 100.0),
                ("kalshi".into(), 0.52, 50.0),
            ],
            vec![],
        );
        let (venue, price) = best_venue(&book, "sell");
        assert_eq!(venue, "kalshi");
        assert!((price - 0.52).abs() < 1e-10);
    }

    #[test]
    fn test_merged_mid_price() {
        let mid = merged_mid_price(
            vec![("A".into(), 0.50, 100.0)],
            vec![("A".into(), 0.52, 100.0)],
        );
        assert!((mid - 0.51).abs() < 1e-10);
    }

    #[test]
    fn test_aggregate_filters_invalid() {
        let book = aggregate_books(
            vec![("A".into(), -0.5, 100.0), ("B".into(), 0.50, 0.0)],
            vec![("A".into(), f64::NAN, 100.0)],
        );
        assert!(book.bids.is_empty());
        assert!(book.asks.is_empty());
    }
}
