from market_agents.exchange.book import LimitOrderBook
from market_agents.exchange.orders import Order, OrderType, Side


def mk(oid, side, qty, price, t=0.0, otype=OrderType.LIMIT):
    return Order(order_id=oid, agent_id=oid, asset=0, side=side, order_type=otype, qty=qty, price=price,
                 submit_time=t, arrival_time=t)


def test_price_time_priority():
    b = LimitOrderBook(0)
    b.rest(mk(1, Side.SELL, 5, 101))
    b.rest(mk(2, Side.SELL, 5, 100))
    b.rest(mk(3, Side.SELL, 5, 100))
    inc = mk(9, Side.BUY, 7, 101)
    fills = b.match(inc)
    assert [(f.resting.order_id, f.qty, f.price) for f in fills] == [(2, 5, 100), (3, 2, 100)]
    assert inc.remaining == 0
    assert b.best_ask() == 100 and b.get(3).remaining == 3
    b.check_invariants()


def test_limit_does_not_cross_beyond_price():
    b = LimitOrderBook(0)
    b.rest(mk(1, Side.SELL, 5, 105))
    inc = mk(2, Side.BUY, 5, 104)
    assert b.match(inc) == [] and inc.remaining == 5
    b.rest(inc)
    assert b.best_bid() == 104 and b.best_ask() == 105
    b.check_invariants()


def test_market_order_walks_book_and_stops_when_empty():
    b = LimitOrderBook(0)
    b.rest(mk(1, Side.BUY, 3, 99))
    b.rest(mk(2, Side.BUY, 3, 98))
    inc = mk(3, Side.SELL, 10, None, otype=OrderType.MARKET)
    fills = b.match(inc)
    assert sum(f.qty for f in fills) == 6 and inc.remaining == 4
    assert b.best_bid() is None
    b.check_invariants()


def test_cancel_and_depth():
    b = LimitOrderBook(0)
    for i in range(1, 6):
        b.rest(mk(i, Side.BUY, i, 100 - i))
    bids, asks = b.depth(3)
    assert bids == [(99, 1), (98, 2), (97, 3)] and asks == []
    assert b.cancel(2).remaining == 2 and 2 not in b
    assert b.cancel(2) is None
    assert b.depth_within(Side.BUY, 2) == 1 + 3
    b.check_invariants()