#!/usr/bin/env python3
"""Rebuild a limit order book from market-by-order (MBO) events and verify it.

Two modes:

* ``--self-test`` (default, no third-party packages): replays a synthetic MBO
  stream against a naive reference implementation and checks every price level
  and queue position after every event.
* ``--mbo FILE [--mbp10 FILE]``: replays a Databento MBO file (``.dbn`` or
  ``.dbn.zst``) and, if an MBP-10 file for the same instrument and time range is
  given, compares the rebuilt top ten levels with the vendor's own aggregation
  after every event. Requires ``pip install databento-dbn zstandard`` (``zstandard``
  is only needed for ``.zst`` files).

Event semantics follow Databento's MBO schema: only A (add), C (cancel),
M (modify) and R (clear) change the book. T (trade) and F (fill) are
informational; the size they remove arrives as a separate C record.
"""

from __future__ import annotations

import argparse
import bisect
import random
import sys
import time
from dataclasses import dataclass

UNDEF_PRICE = 9223372036854775807  # Databento sentinel for "no price"
F_LAST = 128
F_SNAPSHOT = 32


@dataclass(slots=True)
class Order:
    side: str  # "B" or "A"
    price: int  # integer price, 1e-9 units in Databento files
    size: int


class BookSide:
    """One side of the book: price -> ordered queue of resting orders."""

    def __init__(self, is_bid: bool) -> None:
        self.is_bid = is_bid
        self.queues: dict[int, dict[int, int]] = {}  # price -> {order_id: size}, insertion order = queue order
        self._sorted: list[int] = []  # ascending prices; kept in sync with queues

    def add(self, order_id: int, price: int, size: int) -> None:
        queue = self.queues.get(price)
        if queue is None:
            queue = self.queues[price] = {}
            bisect.insort(self._sorted, price)
        queue[order_id] = size

    def remove(self, order_id: int, price: int) -> None:
        queue = self.queues[price]
        del queue[order_id]
        if not queue:
            del self.queues[price]
            index = bisect.bisect_left(self._sorted, price)
            del self._sorted[index]

    def levels(self, depth: int) -> list[tuple[int, int, int]]:
        """Best ``depth`` levels as (price, total size, order count)."""
        prices = self._sorted[::-1] if self.is_bid else self._sorted
        out = []
        for price in prices[:depth]:
            queue = self.queues[price]
            out.append((price, sum(queue.values()), len(queue)))
        return out

    def queue_position(self, order_id: int, price: int) -> tuple[int, int]:
        """(orders ahead, shares ahead) for a resting order."""
        ahead_orders = ahead_size = 0
        for other_id, size in self.queues[price].items():
            if other_id == order_id:
                return ahead_orders, ahead_size
            ahead_orders += 1
            ahead_size += size
        raise KeyError(order_id)


class OrderBook:
    def __init__(self) -> None:
        self.orders: dict[int, Order] = {}
        self.bids = BookSide(is_bid=True)
        self.asks = BookSide(is_bid=False)
        self.anomalies: dict[str, int] = {}

    def _side(self, side: str) -> BookSide:
        return self.bids if side == "B" else self.asks

    def _note(self, kind: str) -> None:
        self.anomalies[kind] = self.anomalies.get(kind, 0) + 1

    def apply(self, action: str, side: str, price: int, size: int, order_id: int) -> None:
        if action == "A":
            if order_id in self.orders:
                self._note("add_duplicate_id")
                self._delete(order_id)
            self.orders[order_id] = Order(side, price, size)
            self._side(side).add(order_id, price, size)
        elif action == "C":
            order = self.orders.get(order_id)
            if order is None:
                self._note("cancel_unknown_id")
                return
            remaining = order.size - size
            if remaining <= 0:
                if remaining < 0:
                    self._note("cancel_exceeds_size")
                self._delete(order_id)
            else:
                order.size = remaining
                self._side(order.side).queues[order.price][order_id] = remaining
        elif action == "M":
            order = self.orders.get(order_id)
            if order is None:
                # Some venues send a modify for an order we never saw; treat as an add.
                self._note("modify_unknown_id_as_add")
                self.apply("A", side, price, size, order_id)
                return
            book_side = self._side(order.side)
            if price != order.price or size > order.size:
                # Price change or size increase loses queue priority: go to the back.
                book_side.remove(order_id, order.price)
                order.price, order.size = price, size
                book_side.add(order_id, price, size)
            else:
                # Size decrease at the same price keeps the position in the queue.
                order.size = size
                book_side.queues[price][order_id] = size
        elif action == "R":
            self.orders.clear()
            self.bids = BookSide(is_bid=True)
            self.asks = BookSide(is_bid=False)
        elif action in ("T", "F", "N"):
            return  # no book change; the matching C record does the work
        else:
            raise ValueError(f"unknown action {action!r}")

    def _delete(self, order_id: int) -> None:
        order = self.orders.pop(order_id)
        self._side(order.side).remove(order_id, order.price)

    def top(self, depth: int = 10) -> tuple[list, list]:
        return self.bids.levels(depth), self.asks.levels(depth)

    def spread(self) -> int | None:
        bids, asks = self.top(1)
        if not bids or not asks:
            return None
        return asks[0][0] - bids[0][0]


# --- self test against a naive reference ------------------------------------


def naive_levels(orders: dict[int, Order], side: str, depth: int) -> list[tuple[int, int, int]]:
    by_price: dict[int, list[int]] = {}
    for order in orders.values():
        if order.side == side:
            by_price.setdefault(order.price, []).append(order.size)
    prices = sorted(by_price, reverse=(side == "B"))[:depth]
    return [(p, sum(by_price[p]), len(by_price[p])) for p in prices]


def self_test(events: int, seed: int) -> None:
    rng = random.Random(seed)
    book = OrderBook()
    reference: dict[int, Order] = {}
    arrival: dict[int, int] = {}  # order_id -> arrival counter, for queue-position checks
    clock = 0
    next_id = 1
    counts = {"A": 0, "C": 0, "M": 0, "R": 0, "T": 0, "F": 0}

    for _ in range(events):
        live = list(reference)
        roll = rng.random()
        if not live or roll < 0.45:
            action = "A"
        elif roll < 0.75:
            action = "C"
        elif roll < 0.95:
            action = "M"
        elif roll < 0.985:
            action = "T"
        else:
            action = "R"
        counts[action] += 1

        if action == "A":
            side = rng.choice("BA")
            price = (100_000 + rng.randint(-20, 20)) * 1_000_000  # $99.98 .. $100.02
            size = rng.randint(1, 500)
            order_id = next_id
            next_id += 1
            book.apply("A", side, price, size, order_id)
            reference[order_id] = Order(side, price, size)
            clock += 1
            arrival[order_id] = clock
        elif action == "C":
            order_id = rng.choice(live)
            ref = reference[order_id]
            size = ref.size if rng.random() < 0.7 else rng.randint(1, ref.size)
            book.apply("C", ref.side, ref.price, size, order_id)
            if size >= ref.size:
                del reference[order_id]
                del arrival[order_id]
            else:
                ref.size -= size
        elif action == "M":
            order_id = rng.choice(live)
            ref = reference[order_id]
            if rng.random() < 0.5:
                price, size = ref.price, rng.randint(1, ref.size)  # shrink, keeps priority
            else:
                price = ref.price + rng.choice([-1, 1]) * 1_000_000
                size = rng.randint(1, 500)
            lost_priority = price != ref.price or size > ref.size
            book.apply("M", ref.side, price, size, order_id)
            ref.price, ref.size = price, size
            if lost_priority:
                clock += 1
                arrival[order_id] = clock
        elif action == "T":
            # A trade arrives as three records: T (aggressor), F (resting order filled),
            # then C (the size actually leaves the book). Only the C changes the reference,
            # so a book that also reduces on F is caught by the level check below.
            order_id = rng.choice(live)
            ref = reference[order_id]
            qty = rng.randint(1, ref.size)
            aggressor = "A" if ref.side == "B" else "B"
            book.apply("T", aggressor, ref.price, qty, 0)
            book.apply("F", ref.side, ref.price, qty, order_id)
            counts["F"] += 1
            book.apply("C", ref.side, ref.price, qty, order_id)
            if qty >= ref.size:
                del reference[order_id]
                del arrival[order_id]
            else:
                ref.size -= qty
        else:
            book.apply("R", "N", UNDEF_PRICE, 0, 0)
            reference.clear()
            arrival.clear()

        # Check every level on both sides.
        for side in "BA":
            got = (book.bids if side == "B" else book.asks).levels(10_000)
            want = naive_levels(reference, side, 10_000)
            if got != want:
                raise AssertionError(f"level mismatch after {action}: {got[:3]} != {want[:3]}")
        # Check queue order at every price: earlier arrival must be ahead.
        for order_id, ref in reference.items():
            ahead_orders, _ = book._side(ref.side).queue_position(order_id, ref.price)
            want_ahead = sum(
                1
                for other_id, other in reference.items()
                if other.side == ref.side and other.price == ref.price and arrival[other_id] < arrival[order_id]
            )
            if ahead_orders != want_ahead:
                raise AssertionError(f"queue position mismatch for order {order_id}")

    assert dict(book.orders) == {k: v for k, v in reference.items()}
    print(f"self-test passed: {events} events {counts}, {len(book.orders)} resting orders, anomalies={book.anomalies}")


# --- Databento replay ---------------------------------------------------------


def replay_databento(mbo_path: str, mbp10_path: str | None, depth: int) -> None:
    from databento_dbn import DBNDecoder  # noqa: PLC0415  (optional dependency)

    def read_bytes(path: str) -> bytes:
        raw = open(path, "rb").read()
        if path.endswith(".zst"):
            import zstandard  # noqa: PLC0415  (installed with databento-dbn's client, or pip install zstandard)

            raw = zstandard.ZstdDecompressor().decompressobj().decompress(raw)
        return raw

    def records(path: str):
        decoder = DBNDecoder()
        decoder.write(read_bytes(path))
        for record in decoder.decode():
            if hasattr(record, "order_id"):
                yield record

    # One book per (publisher, instrument). This script keeps exactly one; a file that
    # mixes several would silently merge their orders, so refuse it.
    mbo_keys: set[tuple[int, int]] = set()
    vendor_keys: set[tuple[int, int]] = set()

    vendor: dict[int, tuple[list, list]] = {}
    if mbp10_path:
        decoder = DBNDecoder()
        decoder.write(read_bytes(mbp10_path))
        for record in decoder.decode():
            if not hasattr(record, "levels"):
                continue
            vendor_keys.add((record.publisher_id, record.instrument_id))
            bids = [(lv.bid_px, lv.bid_sz, lv.bid_ct) for lv in record.levels if lv.bid_px != UNDEF_PRICE]
            asks = [(lv.ask_px, lv.ask_sz, lv.ask_ct) for lv in record.levels if lv.ask_px != UNDEF_PRICE]
            vendor[record.sequence] = (bids, asks)
        if len(vendor_keys) != 1:
            sys.exit(f"MBP-10 file must contain exactly one (publisher, instrument), found {sorted(vendor_keys)}")
        print(f"loaded {len(vendor):,} MBP-10 records for verification, instrument {next(iter(vendor_keys))}")

    book = OrderBook()
    counts: dict[str, int] = {}
    snapshot = 0
    compared = matched = 0
    first_mismatch = None
    started = time.time()
    total = 0
    for record in records(mbo_path):
        total += 1
        mbo_keys.add((record.publisher_id, record.instrument_id))
        if len(mbo_keys) > 1:
            sys.exit(f"MBO file mixes several (publisher, instrument) pairs: {sorted(mbo_keys)}; filter to one first")
        action = str(record.action)
        counts[action] = counts.get(action, 0) + 1
        if record.flags & F_SNAPSHOT:
            snapshot += 1
        book.apply(action, str(record.side), record.price, record.size, record.order_id)
        # Compare only at F_LAST: the venue's event is complete and both books are in a
        # consistent state. Inside an event the vendor's MBP-10 emits the trade record
        # before applying the matching cancel, so mid-event states legitimately differ.
        if vendor and record.flags & F_LAST and record.sequence in vendor:
            compared += 1
            got = book.top(depth)
            if got == vendor[record.sequence]:
                matched += 1
            elif first_mismatch is None:
                first_mismatch = (record.sequence, action, got[0][:2], vendor[record.sequence][0][:2])
    elapsed = time.time() - started
    if total == 0:
        sys.exit("no MBO records found in the input")
    if vendor_keys and vendor_keys != mbo_keys:
        sys.exit(f"MBO instrument {sorted(mbo_keys)} does not match MBP-10 instrument {sorted(vendor_keys)}")

    bids, asks = book.top(5)
    print(f"replayed {total:,} MBO records in {elapsed:.1f}s ({total / elapsed:,.0f} rec/s)")
    print(f"actions: {dict(sorted(counts.items()))}  snapshot records: {snapshot}")
    print(f"resting orders at end: {len(book.orders):,}  anomalies: {book.anomalies}")
    print("best bids:", [(p / 1e9, s, c) for p, s, c in bids])
    print("best asks:", [(p / 1e9, s, c) for p, s, c in asks])
    if vendor:
        print(f"verification at F_LAST events: compared {compared:,}, matched {matched:,}, "
              f"mismatched {compared - matched:,}")
        if compared == 0:
            sys.exit("verification failed: no MBO event matched an MBP-10 sequence number; "
                     "the two files do not cover the same time range")
        if first_mismatch:
            print("first mismatch:", first_mismatch)
        if compared != matched:
            sys.exit(1)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--mbo", help="Databento MBO file (.dbn or .dbn.zst)")
    parser.add_argument("--mbp10", help="Databento MBP-10 file for the same instrument and range")
    parser.add_argument("--depth", type=int, default=10)
    parser.add_argument("--self-test", action="store_true")
    parser.add_argument("--events", type=int, default=20_000)
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if args.mbo:
        replay_databento(args.mbo, args.mbp10, args.depth)
    else:
        self_test(args.events, args.seed)


if __name__ == "__main__":
    main()