pub mod analytics;

use crate::{Quote, Trade, Side, OrderBookSnapshot, PriceLevel};
use chrono::{DateTime, Utc};
use dashmap::DashMap;
use parking_lot::RwLock;
use std::collections::BTreeMap;
use std::sync::Arc;

pub use analytics::{OrderBookAnalytics, LiquidityProfile, OrderBookImbalance, MarketImpact};

// Manages order books for all symbols
pub struct OrderBookManager {
    books: DashMap<String, Arc<RwLock<OrderBook>>>,
}

impl OrderBookManager {
    pub fn new() -> Self {
        Self {
            books: DashMap::new(),
        }
    }

    pub fn get_or_create(&self, symbol: &str) -> Arc<RwLock<OrderBook>> {
        self.books
            .entry(symbol.to_string())
            .or_insert_with(|| Arc::new(RwLock::new(OrderBook::new(symbol.to_string()))))
            .clone()
    }

    pub fn update_quote(&self, symbol: &str, quote: Quote, timestamp: DateTime<Utc>) {
        let book = self.get_or_create(symbol);
        let mut book_guard = book.write();
        book_guard.update_quote(quote, timestamp);
    }

    pub fn update_trade(&self, symbol: &str, trade: Trade, timestamp: DateTime<Utc>) {
        let book = self.get_or_create(symbol);
        let mut book_guard = book.write();
        book_guard.update_trade(trade, timestamp);
    }

    pub fn get_snapshot(&self, symbol: &str, depth: usize) -> Option<OrderBookSnapshot> {
        self.books.get(symbol).map(|book| {
            let book_guard = book.read();
            book_guard.get_snapshot(depth)
        })
    }

    pub fn get_best_bid_ask(&self, symbol: &str) -> Option<(f64, f64)> {
        self.books.get(symbol).and_then(|book| {
            let book_guard = book.read();
            book_guard.get_best_bid_ask()
        })
    }
    
    pub fn get_analytics(&self, symbol: &str, depth: usize) -> Option<OrderBookAnalytics> {
        self.get_snapshot(symbol, depth)
            .and_then(|snapshot| OrderBookAnalytics::calculate(&snapshot))
    }
    
    pub fn get_liquidity_profile(&self, symbol: &str, depth: usize) -> Option<LiquidityProfile> {
        self.get_snapshot(symbol, depth)
            .map(|snapshot| LiquidityProfile::from_snapshot(&snapshot))
    }
}

// Individual order book for a symbol
pub struct OrderBook {
    symbol: String,
    bids: BTreeMap<OrderedFloat, Level>,
    asks: BTreeMap<OrderedFloat, Level>,
    last_update: DateTime<Utc>,
    last_trade_price: Option<f64>,
    last_trade_size: Option<f64>,
}

#[derive(Clone, Debug)]
struct Level {
    price: f64,
    size: f64,
    order_count: u32,
    last_update: DateTime<Utc>,
}

// Wrapper for f64 to allow BTreeMap ordering
#[derive(Clone, Copy, Debug, PartialEq)]
struct OrderedFloat(f64);

impl Eq for OrderedFloat {}

impl PartialOrd for OrderedFloat {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

impl Ord for OrderedFloat {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.partial_cmp(other).unwrap_or(std::cmp::Ordering::Equal)
    }
}

impl OrderBook {
    pub fn new(symbol: String) -> Self {
        Self {
            symbol,
            bids: BTreeMap::new(),
            asks: BTreeMap::new(),
            last_update: Utc::now(),
            last_trade_price: None,
            last_trade_size: None,
        }
    }

    pub fn update_quote(&mut self, quote: Quote, timestamp: DateTime<Utc>) {
        // Update bid
        if quote.bid > 0.0 && quote.bid_size > 0.0 {
            self.bids.insert(
                OrderedFloat(-quote.bid), // Negative for reverse ordering
                Level {
                    price: quote.bid,
                    size: quote.bid_size,
                    order_count: 1,
                    last_update: timestamp,
                },
            );
        }

        // Update ask
        if quote.ask > 0.0 && quote.ask_size > 0.0 {
            self.asks.insert(
                OrderedFloat(quote.ask),
                Level {
                    price: quote.ask,
                    size: quote.ask_size,
                    order_count: 1,
                    last_update: timestamp,
                },
            );
        }

        self.last_update = timestamp;
        self.clean_stale_levels(timestamp);
    }

    pub fn update_trade(&mut self, trade: Trade, timestamp: DateTime<Utc>) {
        self.last_trade_price = Some(trade.price);
        self.last_trade_size = Some(trade.size);
        self.last_update = timestamp;

        // Optionally update order book based on trade
        // Remove liquidity that was likely consumed
        match trade.side {
            Side::Buy => {
                // Trade hit the ask, remove liquidity
                self.remove_liquidity_up_to_asks(trade.price, trade.size);
            }
            Side::Sell => {
                // Trade hit the bid, remove liquidity
                self.remove_liquidity_up_to_bids(trade.price, trade.size);
            }
        }
    }

    pub fn get_snapshot(&self, depth: usize) -> OrderBookSnapshot {
        let bids: Vec<PriceLevel> = self.bids
            .values()
            .take(depth)
            .map(|level| PriceLevel {
                price: level.price,
                size: level.size,
                order_count: Some(level.order_count),
            })
            .collect();

        let asks: Vec<PriceLevel> = self.asks
            .values()
            .take(depth)
            .map(|level| PriceLevel {
                price: level.price,
                size: level.size,
                order_count: Some(level.order_count),
            })
            .collect();

        OrderBookSnapshot {
            symbol: self.symbol.clone(),
            timestamp: self.last_update,
            bids,
            asks,
        }
    }

    pub fn get_best_bid_ask(&self) -> Option<(f64, f64)> {
        let best_bid = self.bids.values().next()?.price;
        let best_ask = self.asks.values().next()?.price;
        Some((best_bid, best_ask))
    }
    
    pub fn get_mid_price(&self) -> Option<f64> {
        self.get_best_bid_ask()
            .map(|(bid, ask)| (bid + ask) / 2.0)
    }
    
    pub fn get_spread(&self) -> Option<f64> {
        self.get_best_bid_ask()
            .map(|(bid, ask)| ask - bid)
    }
    
    pub fn get_depth_at_price(&self, price: f64, side: Side) -> f64 {
        match side {
            Side::Buy => {
                self.bids.values()
                    .filter(|level| level.price >= price)
                    .map(|level| level.size)
                    .sum()
            }
            Side::Sell => {
                self.asks.values()
                    .filter(|level| level.price <= price)
                    .map(|level| level.size)
                    .sum()
            }
        }
    }
    
    pub fn get_volume_weighted_price(&self, size: f64, side: Side) -> Option<f64> {
        let levels: Vec<&Level> = match side {
            Side::Buy => self.asks.values().collect(),
            Side::Sell => self.bids.values().collect(),
        };
        
        let mut remaining_size = size;
        let mut total_cost = 0.0;
        let mut total_shares = 0.0;
        
        for level in levels {
            if remaining_size <= 0.0 {
                break;
            }
            
            let fill_size = remaining_size.min(level.size);
            total_cost += fill_size * level.price;
            total_shares += fill_size;
            remaining_size -= fill_size;
        }
        
        if total_shares > 0.0 {
            Some(total_cost / total_shares)
        } else {
            None
        }
    }

    fn clean_stale_levels(&mut self, current_time: DateTime<Utc>) {
        let stale_threshold = chrono::Duration::seconds(60); // 60 seconds
        
        self.bids.retain(|_, level| {
            current_time - level.last_update < stale_threshold
        });
        
        self.asks.retain(|_, level| {
            current_time - level.last_update < stale_threshold
        });
    }

    fn remove_liquidity_up_to_asks(&mut self, price: f64, size: f64) {
        let mut remaining_size = size;
        let mut to_remove = Vec::new();

        for (key, level) in self.asks.iter_mut() {
            if level.price <= price {
                if level.size <= remaining_size {
                    remaining_size -= level.size;
                    to_remove.push(*key);
                } else {
                    level.size -= remaining_size;
                    break;
                }
            } else {
                break;
            }
        }

        for key in to_remove {
            self.asks.remove(&key);
        }
    }

    fn remove_liquidity_up_to_bids(&mut self, price: f64, size: f64) {
        let mut remaining_size = size;
        let mut to_remove = Vec::new();

        for (key, level) in self.bids.iter_mut() {
            if level.price >= price {
                if level.size <= remaining_size {
                    remaining_size -= level.size;
                    to_remove.push(*key);
                } else {
                    level.size -= remaining_size;
                    break;
                }
            } else {
                break;
            }
        }

        for key in to_remove {
            self.bids.remove(&key);
        }
    }
}