use crate::{
    helpers::get_total_amount,
    market_data::OrderData,
    order_book::Side,
};
use blart::{
    AsBytes,
    NoPrefixesBytes,
    TreeMap,
};
use fuels::{
    tx::TxId,
    types::{
        AssetId,
        Identity,
    },
};
use indexmap::IndexMap;
use slotmap::{
    DefaultKey,
    SlotMap,
};
use std::fmt::Debug;

pub type Price = u64;
pub type Quantity = u64;
pub type OrderId = DefaultKey;
pub type Balances = IndexMap<Identity, (u64, u64)>;

/// big endian representation of price for use as radix key
#[derive(Debug, Clone)]
struct PriceKey([u8; 8]);

impl From<Price> for PriceKey {
    fn from(value: Price) -> Self {
        Self(value.to_be_bytes())
    }
}

impl From<&Price> for PriceKey {
    fn from(value: &Price) -> Self {
        Self(value.to_be_bytes())
    }
}

impl From<&mut Price> for PriceKey {
    fn from(value: &mut Price) -> Self {
        Self(value.to_be_bytes())
    }
}

impl From<PriceKey> for Price {
    fn from(value: PriceKey) -> Self {
        u64::from_be_bytes(value.0)
    }
}

impl From<&PriceKey> for Price {
    fn from(value: &PriceKey) -> Self {
        u64::from_be_bytes(value.0)
    }
}

impl From<&mut PriceKey> for Price {
    fn from(value: &mut PriceKey) -> Self {
        u64::from_be_bytes(value.0)
    }
}

impl AsBytes for PriceKey {
    fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

unsafe impl NoPrefixesBytes for PriceKey {}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Fill {
    pub order_id: u64,
    pub quantity: u64,
    pub price: u64,
}

#[derive(Debug)]
pub struct OrderParams {
    pub side: Side,
    pub price: u64,
    pub quantity: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Order {
    pub trader_id: Identity,
    pub order_id: u64,
    pub side: Side,
    pub price: u64,
    pub quantity: Quantity,
    pub fill: Vec<Fill>,
    pub tx_id: Option<TxId>,
}

impl Order {
    pub fn new(
        order_id: u64,
        OrderData {
            trader_id,
            side,
            price,
            quantity,
        }: OrderData,
    ) -> Self {
        Self {
            order_id,
            trader_id,
            price,
            quantity,
            side,
            fill: vec![],
            tx_id: None,
        }
    }

    pub fn fill(&mut self, order_id: u64, fill_quantity: Quantity, price: Price) -> bool {
        let filled_order = Fill {
            quantity: fill_quantity,
            price,
            order_id,
        };
        self.fill.push(filled_order);
        self.quantity -= fill_quantity;
        self.quantity == 0
    }

    pub fn total(&self, decimals: u64) -> u64 {
        get_total_amount(self.quantity, self.price, decimals)
    }

    pub fn filled_total(&self, decimals: u64) -> u64 {
        self.fill
            .iter()
            .map(|f| get_total_amount(f.quantity, f.price, decimals))
            .sum()
    }

    pub fn remaining_total(&self, decimals: u64) -> Quantity {
        match self.side {
            Side::Buy => {
                get_total_amount(self.total_quantity(), self.price, decimals)
                    - self.filled_total(decimals)
            }
            Side::Sell => 0,
        }
    }

    pub fn total_quantity(&self) -> Quantity {
        self.quantity + self.filled_quantity()
    }

    pub fn filled_quantity(&self) -> Quantity {
        self.fill.iter().map(|q| q.quantity).sum::<Quantity>()
    }

    pub fn filled_price(&self) -> Price {
        let total_price: Price = self.fill.iter().map(|q| q.price).sum();
        let count = self.fill.len();

        if count == 0 {
            0
        } else {
            total_price / count as Price
        }
    }

    pub fn is_closed(&self) -> bool {
        self.quantity == 0
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MatchedEvent {
    pub direction: Side,
    pub taker_id: u64,
    pub maker_id: u64,
    pub quantity: u64,
    pub price: u64,
    pub total: u64,
}

#[derive(Clone, Default)]
pub struct BookConfig {
    pub base_asset: AssetId,
    pub base_decimals: u64,
    pub quote_asset: AssetId,
    pub quote_decimals: u64,
    pub taker_fee: u64,
    pub maker_fee: u64,
}

#[derive(Clone)]
pub struct Book {
    pub order_id: u64,
    pub base_asset: AssetId,
    pub base_decimals: u64,
    pub quote_asset: AssetId,
    pub quote_decimals: u64,
    pub taker_fee: u64,
    pub maker_fee: u64,
    pub all_orders: Vec<Order>,
    pub trades: Vec<MatchedEvent>,
    pub cancels: Vec<Order>,
    pub balances: Balances,
    buys: TreeMap<PriceKey, Vec<OrderId>>,
    sells: TreeMap<PriceKey, Vec<OrderId>>,
    orders: SlotMap<OrderId, Order>,
}

impl Book {
    pub fn new(config: BookConfig) -> Self {
        Self {
            order_id: 0,
            base_asset: config.base_asset,
            base_decimals: config.base_decimals,
            quote_asset: config.quote_asset,
            quote_decimals: config.quote_decimals,
            taker_fee: config.taker_fee,
            maker_fee: config.taker_fee,
            buys: TreeMap::new(),
            sells: TreeMap::new(),
            orders: SlotMap::new(),
            trades: Vec::new(),
            cancels: Vec::new(),
            all_orders: Vec::new(),
            balances: Balances::new(),
        }
    }

    pub fn subtract_fee(&self, is_maker: bool, amount: u64) -> u64 {
        let fee: u64 = if is_maker {
            self.maker_fee
        } else {
            self.taker_fee
        };
        if fee == 0 {
            return amount;
        }
        amount - ((amount * fee) / 1_000_000)
    }

    pub fn subtract_taker_fee(&self, amount: u64) -> u64 {
        self.subtract_fee(false, amount)
    }

    pub fn subtract_maker_fee(&self, amount: u64) -> u64 {
        self.subtract_fee(true, amount)
    }

    pub fn set_address_balance(
        &mut self,
        trader_id: &Identity,
        base_balance: u64,
        quote_balance: u64,
    ) {
        self.balances
            .insert(*trader_id, (base_balance, quote_balance));
    }

    pub fn set_balances(&mut self, balances: Balances) {
        self.balances = balances.clone();
    }

    pub fn change_balance_insert(&mut self, order: &Order) {
        if let Some(balance) = self.balances.get_mut(&order.trader_id) {
            if order.side == Side::Buy {
                balance.1 -= get_total_amount(
                    order.total_quantity(),
                    order.price,
                    self.base_decimals,
                );
            } else {
                balance.0 -= order.total_quantity();
            };
        }
    }

    pub fn change_balance_cancel(&mut self, order: &Order) {
        if let Some(balance) = self.balances.get_mut(&order.trader_id) {
            if order.side == Side::Buy {
                balance.1 +=
                    get_total_amount(order.quantity, order.price, self.base_decimals);
            } else {
                balance.0 += order.quantity;
            };
        }
    }

    pub fn change_balance_match(
        &mut self,
        maker: &Order,
        taker: &Order,
        fill_quantity: u64,
    ) {
        let maker_balances = self.balances.get(&maker.trader_id);
        let taker_balances = self.balances.get(&taker.trader_id);

        if maker_balances.is_none() || taker_balances.is_none() {
            return;
        }
        let mut maker_balances = maker_balances.cloned().unwrap();
        let mut taker_balances = taker_balances.cloned().unwrap();

        if maker.side == Side::Buy {
            maker_balances.0 += self.subtract_maker_fee(fill_quantity);
            taker_balances.1 += self.subtract_taker_fee(get_total_amount(
                fill_quantity,
                maker.price,
                self.base_decimals,
            ));
        } else {
            taker_balances.0 += self.subtract_taker_fee(fill_quantity);
            maker_balances.1 += self.subtract_maker_fee(get_total_amount(
                fill_quantity,
                maker.price,
                self.base_decimals,
            ));
        }

        if taker.is_closed() && taker.remaining_total(self.base_decimals) > 0 {
            taker_balances.1 += taker.remaining_total(self.base_decimals);
        }
        if maker.is_closed() && maker.remaining_total(self.base_decimals) > 0 {
            taker_balances.1 += taker.remaining_total(self.base_decimals);
        }

        self.balances.insert(maker.trader_id, maker_balances);
        self.balances.insert(taker.trader_id, taker_balances);
    }

    pub fn load(&mut self, orders: Vec<Order>) {
        for order in orders {
            if order.quantity == 0 {
                continue;
            }
            self.insert_order(&order);
        }
    }

    pub fn get_buys_count(&mut self) -> usize {
        self.buys.len()
    }

    pub fn get_sells_count(&mut self) -> usize {
        self.sells.len()
    }

    fn insert_order(&mut self, order: &Order) -> OrderId {
        let price = order.price;
        let id = self.orders.insert(order.clone());
        self.change_balance_insert(order);

        match order.side {
            Side::Sell => {
                self.sells.entry(price.into()).or_default().push(id);
            }
            Side::Buy => {
                self.buys.entry(price.into()).or_default().push(id);
            }
        }
        self.execute();
        id
    }

    pub fn get_orders(&self) -> Vec<Order> {
        self.orders.values().cloned().collect()
    }

    pub fn cancel(&mut self, order_id: u64) {
        if let Some((order_key, order)) =
            self.orders.iter().find(|o| o.1.order_id == order_id)
        {
            self.cancels.push(order.clone());
            self.change_balance_cancel(&order.clone());
            self.remove_order(order_key);
        }
    }

    fn remove_order(&mut self, order_id: OrderId) {
        let order = &self.orders[order_id];
        match order.side {
            Side::Sell => {
                if let Some(ref_vec) = self.sells.get_mut(&order.price.into()) {
                    ref_vec.retain(|id| *id != order_id);
                    if ref_vec.is_empty() {
                        self.sells.remove(&order.price.into());
                    }
                }
            }
            Side::Buy => {
                if let Some(ref_vec) = self.buys.get_mut(&order.price.into()) {
                    ref_vec.retain(|id| *id != order_id);
                    if ref_vec.is_empty() {
                        self.buys.remove(&order.price.into());
                    }
                }
            }
        }
        self.orders.remove(order_id);
    }

    fn lowest_sell(&self) -> Option<(&PriceKey, &OrderId)> {
        match self.sells.first_key_value().map(|(k, v)| (k, v.first())) {
            Some((k, Some(v))) => Some((k, v)),
            _ => None,
        }
    }

    fn highest_buy(&self) -> Option<(&PriceKey, &OrderId)> {
        match self.buys.last_key_value().map(|(k, v)| (k, v.first())) {
            Some((k, Some(v))) => Some((k, v)),
            _ => None,
        }
    }

    pub fn create_order_matched_event(&self, maker_order: &Order) -> MatchedEvent {
        let last_filled_order = maker_order.fill.last().unwrap();
        let total = get_total_amount(
            last_filled_order.quantity,
            last_filled_order.price,
            self.base_decimals,
        );
        MatchedEvent {
            direction: maker_order.side,
            taker_id: last_filled_order.order_id,
            maker_id: maker_order.order_id,
            quantity: last_filled_order.quantity,
            price: last_filled_order.price,
            total,
        }
    }

    pub fn execute(&mut self) {
        while self
            .lowest_sell()
            .zip(self.highest_buy())
            .map(|((sell, _), (buy, _))| Price::from(buy) >= Price::from(sell))
            .unwrap_or(false)
        {
            let sell_id = *self.lowest_sell().unwrap().1;
            let sell_quantity = self.orders[sell_id].quantity;
            let sell_order_id = self.orders[sell_id].order_id;
            let buy_id = *self.highest_buy().unwrap().1;
            let buy_quantity = self.orders[buy_id].quantity;
            let buy_order_id = self.orders[buy_id].order_id;
            let fill_quantity = sell_quantity.min(buy_quantity);

            let maker_id = if sell_order_id < buy_order_id {
                sell_id
            } else {
                buy_id
            };
            let taker_id = if sell_order_id > buy_order_id {
                sell_id
            } else {
                buy_id
            };
            let maker_price = self.orders[maker_id].price;

            let sell_is_full =
                self.orders[sell_id].fill(buy_order_id, fill_quantity, maker_price);
            let buy_is_full =
                self.orders[buy_id].fill(sell_order_id, fill_quantity, maker_price);

            self.change_balance_match(
                &self.orders[maker_id].clone(),
                &self.orders[taker_id].clone(),
                fill_quantity,
            );
            self.trades
                .push(self.create_order_matched_event(&self.orders[maker_id]));

            if sell_is_full {
                self.remove_order(sell_id);
            }
            if buy_is_full {
                self.remove_order(buy_id);
            }
        }
    }

    pub fn create_order(&mut self, order_data: OrderData) -> Order {
        self.order_id += 1;
        let order = Order::new(self.order_id, order_data);
        self.insert_order(&order);
        self.all_orders.push(order.clone());
        order
    }

    pub fn create_orders(&mut self, orders_data: Vec<OrderData>) -> Vec<Order> {
        let mut orders = Vec::new();
        for order_data in orders_data {
            orders.push(self.create_order(order_data));
        }
        orders
    }
}

#[derive(Clone, Debug)]
pub struct TestOrderBookState {
    pub orders: Vec<OrderData>,
    pub matches: Vec<MatchedEvent>,
    pub cancels: Vec<Order>,
}

pub fn get_default_orders_test(
    base_asset: AssetId,
    quote_asset: AssetId,
    initial_balances: Balances,
) -> (Book, TestOrderBookState) {
    let mut book: Book = Book::new(BookConfig {
        base_asset,
        base_decimals: 9,
        quote_asset,
        quote_decimals: 6,
        taker_fee: 0,
        maker_fee: 0,
    });
    let traders = initial_balances.keys().cloned().collect::<Vec<_>>();

    // Set initial balances
    book.set_balances(initial_balances);

    let mut orders = vec![
        OrderData {
            side: Side::Buy,
            price: 9_500_000,
            quantity: 20_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Sell,
            price: 20_000_000,
            quantity: 20_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Sell,
            price: 10_000_000,
            quantity: 20_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Buy,
            price: 10_000_000,
            quantity: 40_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Sell,
            price: 9_900_000,
            quantity: 40_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Buy,
            price: 10_000_000,
            quantity: 20_000_000,
            trader_id: Identity::default(),
        },
    ];

    for (i, order) in orders.iter_mut().enumerate() {
        let trader_id = traders.get(i % traders.len()).unwrap();
        order.trader_id = *trader_id;
    }

    let matches = vec![
        MatchedEvent {
            direction: Side::Sell,
            maker_id: 3,
            taker_id: 4,
            quantity: 20_000_000,
            price: 10_000_000,
            total: 200_000,
        },
        MatchedEvent {
            direction: Side::Buy,
            maker_id: 4,
            taker_id: 5,
            quantity: 20_000_000,
            price: 10_000_000,
            total: 200_000,
        },
        MatchedEvent {
            direction: Side::Sell,
            maker_id: 5,
            taker_id: 6,
            quantity: 20_000_000,
            price: 9_900_000,
            total: 198_000,
        },
    ];
    let cancels = vec![Order::new(
        1,
        OrderData {
            side: Side::Buy,
            price: 9_500_000,
            quantity: 20_000_000,
            trader_id: *traders.first().unwrap(),
        },
    )];

    // Create book with states
    book.create_orders(orders.clone());

    let test_order_book_state = TestOrderBookState {
        orders,
        matches,
        cancels,
    };

    (book, test_order_book_state)
}

pub fn get_default_orders_with_fees_test(
    base_asset: AssetId,
    quote_asset: AssetId,
    initial_balances: Balances,
) -> (Book, TestOrderBookState) {
    let mut book: Book = Book::new(BookConfig {
        base_asset,
        base_decimals: 9,
        quote_asset,
        quote_decimals: 6,
        taker_fee: 2_000,
        maker_fee: 1_000,
    });
    let traders = initial_balances.keys().cloned().collect::<Vec<_>>();

    // Set initial balances
    book.set_balances(initial_balances);

    let mut orders = vec![
        OrderData {
            side: Side::Buy,
            price: 9_500_000,
            quantity: 20_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Sell,
            price: 20_000_000,
            quantity: 20_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Sell,
            price: 10_000_000,
            quantity: 20_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Buy,
            price: 10_000_000,
            quantity: 40_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Sell,
            price: 9_900_000,
            quantity: 40_000_000,
            trader_id: Identity::default(),
        },
        OrderData {
            side: Side::Buy,
            price: 10_000_000,
            quantity: 20_000_000,
            trader_id: Identity::default(),
        },
    ];

    for (i, order) in orders.iter_mut().enumerate() {
        let trader_id = traders.get(i % traders.len()).unwrap();
        order.trader_id = *trader_id;
    }

    let matches = vec![
        MatchedEvent {
            direction: Side::Sell,
            maker_id: 3,
            taker_id: 4,
            quantity: 20_000_000,
            price: 10_000_000,
            total: 200_000,
        },
        MatchedEvent {
            direction: Side::Buy,
            maker_id: 4,
            taker_id: 5,
            quantity: 20_000_000,
            price: 10_000_000,
            total: 200_000,
        },
        MatchedEvent {
            direction: Side::Sell,
            maker_id: 5,
            taker_id: 6,
            quantity: 20_000_000,
            price: 9_900_000,
            total: 198_000,
        },
    ];
    let cancels = vec![Order::new(
        1,
        OrderData {
            side: Side::Buy,
            price: 9_500_000,
            quantity: 20_000_000,
            trader_id: *traders.first().unwrap(),
        },
    )];

    // Create book with states
    book.create_orders(orders.clone());

    let test_order_book_state = TestOrderBookState {
        orders,
        matches,
        cancels,
    };

    (book, test_order_book_state)
}

#[tokio::test]
async fn test_order_book_data() {
    use fuels::types::{
        Address,
        Identity,
    };
    use std::str::FromStr;
    const USDC: &str =
        "0x8288f40f7c9ab3b5bc36783b692f4ae96975c5a04d8f0da258886fdada0b26af";
    const BTC: &str =
        "0xd0cef4b9cb0706cfa4aa174e75fb006b8ddd2a0ecd6296a27c0a05dcb856f32f";
    let address_1 = Identity::Address(Address::from([5u8; 32]));
    let address_2 = Identity::Address(Address::from([6u8; 32]));

    let balances = Balances::from_iter(vec![
        (address_1, (100_000_000_000_000, 100_000_000_000_000)),
        (address_2, (100_000_000_000_000, 100_000_000_000_000)),
    ]);

    let (mut book, test_order_book_state) = get_default_orders_test(
        AssetId::from_str(BTC).unwrap(),
        AssetId::from_str(USDC).unwrap(),
        balances,
    );
    book.cancel(1);
    let balances = Balances::from_iter(vec![
        (address_1, (99999940000000, 100000000598000)),
        (address_2, (100000040000000, 99999999402000)),
    ]);

    assert_eq!(test_order_book_state.matches, book.trades);
    assert_eq!(test_order_book_state.cancels, book.cancels);
    assert_eq!(balances, book.balances);
}

#[tokio::test]
async fn test_order_book_data_fees() {
    use fuels::types::{
        Address,
        Identity,
    };
    use std::str::FromStr;
    const USDC: &str =
        "0x8288f40f7c9ab3b5bc36783b692f4ae96975c5a04d8f0da258886fdada0b26af";
    const BTC: &str =
        "0xd0cef4b9cb0706cfa4aa174e75fb006b8ddd2a0ecd6296a27c0a05dcb856f32f";
    let address_1 = Identity::Address(Address::from([5u8; 32]));
    let address_2 = Identity::Address(Address::from([6u8; 32]));

    let balances = Balances::from_iter(vec![
        (address_1, (100_000_000_000_000, 100_000_000_000_000)),
        (address_2, (100_000_000_000_000, 100_000_000_000_000)),
    ]);

    let (mut book, test_order_book_state) = get_default_orders_with_fees_test(
        AssetId::from_str(BTC).unwrap(),
        AssetId::from_str(USDC).unwrap(),
        balances,
    );
    book.cancel(1);
    let balances = Balances::from_iter(vec![
        (address_1, (99999940000000, 100000000596804)),
        (address_2, (100000039880000, 99999999402000)),
    ]);

    assert_eq!(test_order_book_state.matches, book.trades);
    assert_eq!(test_order_book_state.cancels, book.cancels);
    assert_eq!(balances, book.balances);
}
