nautilus_databento/

loader.rs

// -------------------------------------------------------------------------------------------------
//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
//  https://nautechsystems.io
//
//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
//  You may not use this file except in compliance with the License.
//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// -------------------------------------------------------------------------------------------------

use std::{
    env,
    path::{Path, PathBuf},
};

use ahash::AHashMap;
use anyhow::Context;
use databento::dbn::{self, InstrumentDefMsg};
use dbn::{
    Publisher,
    decode::{DbnMetadata, DecodeStream, dbn::Decoder},
};
use fallible_streaming_iterator::FallibleStreamingIterator;
use indexmap::IndexMap;
use nautilus_model::{
    data::{Bar, Data, InstrumentStatus, OrderBookDelta, OrderBookDepth10, QuoteTick, TradeTick},
    identifiers::{InstrumentId, Symbol, Venue},
    instruments::{Instrument, InstrumentAny},
};

use super::{
    decode::{
        MboDeltaBuffer, decode_imbalance_msg, decode_mbo_msg, decode_record, decode_statistics_msg,
        decode_status_msg, is_supported_stat_type,
    },
    symbology::decode_nautilus_instrument_id,
    types::{DatabentoImbalance, DatabentoPublisher, DatabentoStatistics, Dataset, PublisherId},
};
use crate::{
    common::{build_publisher_venue_map, load_publishers},
    decode::{DatabentoDecodeConfig, decode_instrument_def_msg},
    symbology::MetadataCache,
};

/// A Nautilus data loader for Databento Binary Encoding (DBN) format data.
///
/// # Supported Schemas
///  - `MBO` -> `OrderBookDelta`
///  - `MBP_1` -> `(QuoteTick, Option<TradeTick>)`
///  - `MBP_10` -> `OrderBookDepth10`
///  - `BBO_1S` -> `QuoteTick`
///  - `BBO_1M` -> `QuoteTick`
///  - `CMBP_1` -> `(QuoteTick, Option<TradeTick>)`
///  - `CBBO_1S` -> `QuoteTick`
///  - `CBBO_1M` -> `QuoteTick`
///  - `TCBBO` -> `(QuoteTick, TradeTick)`
///  - `TBBO` -> `(QuoteTick, TradeTick)`
///  - `TRADES` -> `TradeTick`
///  - `OHLCV_1S` -> `Bar`
///  - `OHLCV_1M` -> `Bar`
///  - `OHLCV_1H` -> `Bar`
///  - `OHLCV_1D` -> `Bar`
///  - `OHLCV_EOD` -> `Bar`
///  - `DEFINITION` -> `Instrument`
///  - `IMBALANCE` -> `DatabentoImbalance`
///  - `STATISTICS` -> `DatabentoStatistics`
///  - `STATUS` -> `InstrumentStatus`
///
/// # References
///
/// <https://databento.com/docs/schemas-and-data-formats>
#[cfg_attr(
    feature = "python",
    pyo3::pyclass(module = "nautilus_trader.adapters.databento")
)]
#[cfg_attr(
    feature = "python",
    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.adapters.databento")
)]
#[derive(Debug)]
pub struct DatabentoDataLoader {
    publishers_map: IndexMap<PublisherId, DatabentoPublisher>,
    venue_dataset_map: IndexMap<Venue, Dataset>,
    publisher_venue_map: IndexMap<PublisherId, Venue>,
    symbol_venue_map: AHashMap<Symbol, Venue>,
    price_precisions: AHashMap<Symbol, u8>,
}

impl DatabentoDataLoader {
    /// Creates a new [`DatabentoDataLoader`] instance.
    ///
    /// # Errors
    ///
    /// Returns an error if locating or loading publishers data fails.
    pub fn new(publishers_filepath: Option<PathBuf>) -> anyhow::Result<Self> {
        let mut loader = Self {
            publishers_map: IndexMap::new(),
            venue_dataset_map: IndexMap::new(),
            publisher_venue_map: IndexMap::new(),
            symbol_venue_map: AHashMap::new(),
            price_precisions: AHashMap::new(),
        };

        // Load publishers
        let publishers_filepath = if let Some(p) = publishers_filepath {
            p
        } else {
            // Use built-in publishers path
            let mut exe_path = env::current_exe()?;
            exe_path.pop();
            exe_path.push("publishers.json");
            exe_path
        };

        loader
            .load_publishers(publishers_filepath)
            .context("error loading publishers.json")?;

        Ok(loader)
    }

    /// Load the publishers data from the file at the given `filepath`.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be read or parsed as JSON.
    pub fn load_publishers(&mut self, filepath: PathBuf) -> anyhow::Result<()> {
        let publishers = load_publishers(filepath)?;

        self.publishers_map = publishers
            .iter()
            .cloned()
            .map(|p| (p.publisher_id, p))
            .collect();

        let mut venue_dataset_map = IndexMap::new();

        // Only insert a dataset if the venue key is not already in the map
        for publisher in &publishers {
            let venue = Venue::from(publisher.venue.as_str());
            let dataset = Dataset::from(publisher.dataset.as_str());
            venue_dataset_map.entry(venue).or_insert(dataset);
        }

        self.venue_dataset_map = venue_dataset_map;
        apply_default_venue_dataset_mappings(&mut self.venue_dataset_map);

        self.publisher_venue_map = build_publisher_venue_map(&publishers);

        Ok(())
    }

    /// Returns the internal Databento publishers currently held by the loader.
    #[must_use]
    pub const fn get_publishers(&self) -> &IndexMap<u16, DatabentoPublisher> {
        &self.publishers_map
    }

    /// Sets the `venue` to map to the given `dataset`.
    pub fn set_dataset_for_venue(&mut self, dataset: Dataset, venue: Venue) {
        _ = self.venue_dataset_map.insert(venue, dataset);
    }

    /// Returns the dataset which matches the given `venue` (if found).
    #[must_use]
    pub fn get_dataset_for_venue(&self, venue: &Venue) -> Option<&Dataset> {
        self.venue_dataset_map.get(venue)
    }

    /// Returns the venue which matches the given `publisher_id` (if found).
    #[must_use]
    pub fn get_venue_for_publisher(&self, publisher_id: PublisherId) -> Option<&Venue> {
        self.publisher_venue_map.get(&publisher_id)
    }

    /// Caches a `price_precision` for the given `symbol`.
    ///
    /// When market data is read without an explicit `price_precision` argument,
    /// the loader resolves precision per record from this cache. Definitions
    /// loaded via [`Self::load_instruments`] are inserted automatically.
    pub fn set_price_precision(&mut self, symbol: Symbol, price_precision: u8) {
        self.price_precisions.insert(symbol, price_precision);
    }

    /// Returns the cached price precisions keyed by symbol.
    #[must_use]
    pub const fn get_price_precisions(&self) -> &AHashMap<Symbol, u8> {
        &self.price_precisions
    }

    /// Resolves a price precision for the given `instrument_id`.
    ///
    /// Resolution order:
    /// 1. The explicit `price_precision` argument (if `Some`).
    /// 2. The cached precision for the instrument's symbol.
    ///
    /// # Errors
    ///
    /// Returns an error when no precision is available.
    fn resolve_price_precision(
        &self,
        instrument_id: &InstrumentId,
        price_precision: Option<u8>,
    ) -> anyhow::Result<u8> {
        if let Some(precision) = price_precision {
            return Ok(precision);
        }

        self.price_precisions
            .get(&instrument_id.symbol)
            .copied()
            .ok_or_else(|| {
                anyhow::anyhow!(
                    "Could not resolve `price_precision` for {instrument_id}: \
                     pass `price_precision` explicitly, call `set_price_precision`, \
                     or load the instrument definitions first via `load_instruments`"
                )
            })
    }

    /// Returns the schema for the given `filepath`.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be decoded or metadata retrieval fails.
    pub fn schema_from_file(&self, filepath: &Path) -> anyhow::Result<Option<String>> {
        let decoder = Decoder::from_zstd_file(filepath)?;
        let metadata = decoder.metadata();
        Ok(metadata.schema.map(|schema| schema.to_string()))
    }

    /// Reads instrument definition records from a DBN file.
    ///
    /// # Errors
    ///
    /// Returns an error if decoding the definition records fails.
    pub fn read_definition_records<'a>(
        &'a mut self,
        filepath: &Path,
        use_exchange_as_venue: bool,
        decode_config: Option<&'a DatabentoDecodeConfig>,
    ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<InstrumentAny>> + 'a> {
        let decoder = Decoder::from_zstd_file(filepath)?;
        let mut dbn_stream = decoder.decode_stream::<InstrumentDefMsg>();

        // Loop over skipped records (Ok(None)) so one unsupported class does not
        // terminate the stream
        Ok(std::iter::from_fn(move || {
            loop {
                let advance = dbn_stream
                    .advance()
                    .map_err(|e| anyhow::anyhow!("Stream advance error: {e}"));
                if let Err(e) = advance {
                    return Some(Err(e));
                }

                let rec = dbn_stream.get()?;

                let result: anyhow::Result<Option<InstrumentAny>> = (|| {
                    let record = dbn::RecordRef::from(rec);
                    let msg = record
                        .get::<InstrumentDefMsg>()
                        .ok_or_else(|| anyhow::anyhow!("Failed to decode InstrumentDefMsg"))?;

                    let raw_symbol = rec
                        .raw_symbol()
                        .map_err(|e| anyhow::anyhow!("Error decoding `raw_symbol`: {e}"))?;
                    let symbol = Symbol::from(raw_symbol);

                    let publisher = rec
                        .hd
                        .publisher()
                        .map_err(|e| anyhow::anyhow!("Invalid `publisher` for record: {e}"))?;
                    let venue = match publisher {
                        Publisher::GlbxMdp3Glbx if use_exchange_as_venue => {
                            let exchange = rec.exchange().map_err(|e| {
                                anyhow::anyhow!("Missing `exchange` for record: {e}")
                            })?;
                            let venue = Venue::from_code(exchange).map_err(|e| {
                                anyhow::anyhow!("Venue not found for exchange {exchange}: {e}")
                            })?;
                            self.symbol_venue_map.insert(symbol, venue);
                            venue
                        }
                        _ => *self
                            .publisher_venue_map
                            .get(&msg.hd.publisher_id)
                            .ok_or_else(|| {
                                anyhow::anyhow!(
                                    "Venue not found for publisher_id {}",
                                    msg.hd.publisher_id
                                )
                            })?,
                    };
                    let instrument_id = InstrumentId::new(symbol, venue);
                    let ts_init = msg.ts_recv.into();

                    decode_instrument_def_msg(rec, instrument_id, Some(ts_init), decode_config)
                })();

                match result {
                    Ok(Some(item)) => return Some(Ok(item)),
                    Ok(None) => {}
                    Err(e) => return Some(Err(e)),
                }
            }
        }))
    }

    /// Reads and decodes market data records from a DBN file.
    ///
    /// # Errors
    ///
    /// Returns an error if reading records fails.
    pub fn read_records<T>(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
        include_trades: bool,
        bars_timestamp_on_close: Option<bool>,
    ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<(Option<Data>, Option<Data>)>> + '_>
    where
        T: dbn::Record + dbn::HasRType + 'static,
    {
        let decoder = Decoder::from_zstd_file(filepath)?;
        let mut metadata_cache = if instrument_id.is_none() {
            Some(MetadataCache::new(decoder.metadata().clone()))
        } else {
            None
        };
        let mut dbn_stream = decoder.decode_stream::<T>();
        let fixed_instrument_id = instrument_id.is_some();
        let mut fixed_price_precision = price_precision;

        Ok(std::iter::from_fn(move || {
            let result: anyhow::Result<Option<(Option<Data>, Option<Data>)>> = (|| {
                dbn_stream
                    .advance()
                    .map_err(|e| anyhow::anyhow!("Stream advance error: {e}"))?;

                if let Some(rec) = dbn_stream.get() {
                    let record = dbn::RecordRef::from(rec);
                    let instrument_id = self
                        .resolve_record_instrument_id(&record, instrument_id, &mut metadata_cache)
                        .context("failed to decode instrument id")?;
                    let resolved_precision = self.resolve_stream_price_precision(
                        &instrument_id,
                        fixed_instrument_id,
                        &mut fixed_price_precision,
                    )?;
                    let (item1, item2) = decode_record(
                        &record,
                        instrument_id,
                        resolved_precision,
                        None,
                        include_trades,
                        bars_timestamp_on_close.unwrap_or(true),
                    )?;
                    Ok(Some((item1, item2)))
                } else {
                    Ok(None)
                }
            })();

            match result {
                Ok(Some(v)) => Some(Ok(v)),
                Ok(None) => None,
                Err(e) => Some(Err(e)),
            }
        }))
    }

    /// Loads all instrument definitions from a DBN file.
    ///
    /// When `skip_on_error` is true, instruments that fail to decode are logged
    /// as warnings and skipped. When false (default), any decode error is propagated.
    ///
    /// # Errors
    ///
    /// Returns an error if loading instruments fails.
    pub fn load_instruments(
        &mut self,
        filepath: &Path,
        use_exchange_as_venue: bool,
        skip_on_error: bool,
        decode_config: Option<&DatabentoDecodeConfig>,
    ) -> anyhow::Result<Vec<InstrumentAny>> {
        let instruments = if skip_on_error {
            let mut collected = Vec::new();

            for result in
                self.read_definition_records(filepath, use_exchange_as_venue, decode_config)?
            {
                match result {
                    Ok(instrument) => collected.push(instrument),
                    Err(e) => log::warn!("Skipping instrument: {e}"),
                }
            }
            collected
        } else {
            self.read_definition_records(filepath, use_exchange_as_venue, decode_config)?
                .collect::<Result<Vec<_>, _>>()?
        };

        for instrument in &instruments {
            self.price_precisions
                .insert(instrument.id().symbol, instrument.price_precision());
        }

        Ok(instruments)
    }

    /// Loads order book delta messages from a DBN MBO schema file.
    ///
    /// Cannot include trades.
    ///
    /// # Errors
    ///
    /// Returns an error if loading order book deltas fails.
    pub fn load_order_book_deltas(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<Vec<OrderBookDelta>> {
        self.read_order_book_deltas(filepath, instrument_id, price_precision)?
            .collect()
    }

    /// Reads order book delta messages from a DBN MBO schema file without collecting them.
    ///
    /// Cannot include trades.
    ///
    /// # Errors
    ///
    /// Returns an error if opening or decoding order book deltas fails.
    pub fn read_order_book_deltas(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<OrderBookDelta>> + '_> {
        let decoder = Decoder::from_zstd_file(filepath)?;
        let mut metadata_cache = if instrument_id.is_none() {
            Some(MetadataCache::new(decoder.metadata().clone()))
        } else {
            None
        };
        let mut dbn_stream = decoder.decode_stream::<dbn::MboMsg>();
        let fixed_instrument_id = instrument_id.is_some();
        let mut fixed_price_precision = price_precision;
        let mut delta_buffer = MboDeltaBuffer::default();
        let mut terminal_error = None;
        let mut finished = false;

        Ok(std::iter::from_fn(move || {
            loop {
                if let Some(delta) = delta_buffer.pop_ready() {
                    return Some(Ok(delta));
                }

                if finished {
                    return terminal_error.take().map(Err);
                }

                let result: anyhow::Result<bool> = (|| {
                    dbn_stream
                        .advance()
                        .map_err(|e| anyhow::anyhow!("Stream advance error: {e}"))?;

                    let Some(rec) = dbn_stream.get() else {
                        return Ok(false);
                    };
                    let record = dbn::RecordRef::from(rec);
                    let instrument_id = self
                        .resolve_record_instrument_id(&record, instrument_id, &mut metadata_cache)
                        .context("failed to decode instrument id")?;
                    let resolved_precision = self.resolve_stream_price_precision(
                        &instrument_id,
                        fixed_instrument_id,
                        &mut fixed_price_precision,
                    )?;
                    let msg = record
                        .get::<dbn::MboMsg>()
                        .ok_or_else(|| anyhow::anyhow!("Failed to decode MboMsg"))?;
                    let (delta, _trade) =
                        decode_mbo_msg(msg, instrument_id, resolved_precision, None, false)?;
                    delta_buffer.push(msg, instrument_id, delta);
                    Ok(true)
                })();

                match result {
                    Ok(true) => {}
                    Ok(false) => {
                        delta_buffer.finish();
                        finished = true;
                    }
                    Err(e) => {
                        delta_buffer.finish();
                        terminal_error = Some(e);
                        finished = true;
                    }
                }
            }
        }))
    }

    /// Loads order book depth10 snapshots from a DBN MBP-10 schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading order book depth10 fails.
    pub fn load_order_book_depth10(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<Vec<OrderBookDepth10>> {
        self.read_records::<dbn::Mbp10Msg>(filepath, instrument_id, price_precision, false, None)?
            .filter_map(|result| match result {
                Ok((Some(item1), _)) => {
                    if let Data::Depth10(depth) = item1 {
                        Some(Ok(*depth))
                    } else {
                        None
                    }
                }
                Ok((None, _)) => None,
                Err(e) => Some(Err(e)),
            })
            .collect()
    }

    /// Loads quote tick messages from a DBN MBP-1 or TBBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading quotes fails.
    pub fn load_quotes(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<Vec<QuoteTick>> {
        self.read_records::<dbn::Mbp1Msg>(filepath, instrument_id, price_precision, false, None)?
            .filter_map(|result| match result {
                Ok((Some(item1), _)) => {
                    if let Data::Quote(quote) = item1 {
                        Some(Ok(quote))
                    } else {
                        None
                    }
                }
                Ok((None, _)) => None,
                Err(e) => Some(Err(e)),
            })
            .collect()
    }

    /// Loads best bid/offer quote messages from a DBN BBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading BBO quotes fails.
    pub fn load_bbo_quotes(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<Vec<QuoteTick>> {
        self.read_records::<dbn::BboMsg>(filepath, instrument_id, price_precision, false, None)?
            .filter_map(|result| match result {
                Ok((Some(item1), _)) => {
                    if let Data::Quote(quote) = item1 {
                        Some(Ok(quote))
                    } else {
                        None
                    }
                }
                Ok((None, _)) => None,
                Err(e) => Some(Err(e)),
            })
            .collect()
    }

    /// Loads consolidated MBP-1 quote messages from a DBN CMBP-1 schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading consolidated MBP-1 quotes fails.
    pub fn load_cmbp_quotes(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<Vec<QuoteTick>> {
        self.read_records::<dbn::Cmbp1Msg>(filepath, instrument_id, price_precision, false, None)?
            .filter_map(|result| match result {
                Ok((Some(item1), _)) => {
                    if let Data::Quote(quote) = item1 {
                        Some(Ok(quote))
                    } else {
                        None
                    }
                }
                Ok((None, _)) => None,
                Err(e) => Some(Err(e)),
            })
            .collect()
    }

    /// Loads consolidated best bid/offer quote messages from a DBN CBBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading consolidated BBO quotes fails.
    pub fn load_cbbo_quotes(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<Vec<QuoteTick>> {
        self.read_records::<dbn::CbboMsg>(filepath, instrument_id, price_precision, false, None)?
            .filter_map(|result| match result {
                Ok((Some(item1), _)) => {
                    if let Data::Quote(quote) = item1 {
                        Some(Ok(quote))
                    } else {
                        None
                    }
                }
                Ok((None, _)) => None,
                Err(e) => Some(Err(e)),
            })
            .collect()
    }

    /// Loads trade messages from a DBN TBBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading TBBO trades fails.
    pub fn load_tbbo_trades(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<Vec<TradeTick>> {
        self.read_records::<dbn::TbboMsg>(filepath, instrument_id, price_precision, true, None)?
            .filter_map(|result| match result {
                Ok((_, maybe_item2)) => {
                    if let Some(Data::Trade(trade)) = maybe_item2 {
                        Some(Ok(trade))
                    } else {
                        None
                    }
                }
                Err(e) => Some(Err(e)),
            })
            .collect()
    }

    /// Loads trade messages from a DBN TCBBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading TCBBO trades fails.
    pub fn load_tcbbo_trades(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<Vec<TradeTick>> {
        self.read_records::<dbn::TcbboMsg>(filepath, instrument_id, price_precision, true, None)?
            .filter_map(|result| match result {
                Ok((_, maybe_item2)) => {
                    if let Some(Data::Trade(trade)) = maybe_item2 {
                        Some(Ok(trade))
                    } else {
                        None
                    }
                }
                Err(e) => Some(Err(e)),
            })
            .collect()
    }

    /// Loads trade messages from a DBN TRADES schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading trades fails.
    pub fn load_trades(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<Vec<TradeTick>> {
        self.read_records::<dbn::TradeMsg>(filepath, instrument_id, price_precision, false, None)?
            .filter_map(|result| match result {
                Ok((Some(item1), _)) => {
                    if let Data::Trade(trade) = item1 {
                        Some(Ok(trade))
                    } else {
                        None
                    }
                }
                Ok((None, _)) => None,
                Err(e) => Some(Err(e)),
            })
            .collect()
    }

    /// Loads OHLCV bar messages from a DBN OHLCV schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading bars fails.
    pub fn load_bars(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
        timestamp_on_close: Option<bool>,
    ) -> anyhow::Result<Vec<Bar>> {
        self.read_records::<dbn::OhlcvMsg>(
            filepath,
            instrument_id,
            price_precision,
            false,
            timestamp_on_close,
        )?
        .filter_map(|result| match result {
            Ok((Some(item1), _)) => {
                if let Data::Bar(bar) = item1 {
                    Some(Ok(bar))
                } else {
                    None
                }
            }
            Ok((None, _)) => None,
            Err(e) => Some(Err(e)),
        })
        .collect()
    }

    /// Loads instrument status messages from a DBN STATUS schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading status records fails.
    pub fn load_status_records<T>(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
    ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<InstrumentStatus>> + '_>
    where
        T: dbn::Record + dbn::HasRType + 'static,
    {
        let decoder = Decoder::from_zstd_file(filepath)?;
        let mut metadata_cache = if instrument_id.is_none() {
            Some(MetadataCache::new(decoder.metadata().clone()))
        } else {
            None
        };
        let mut dbn_stream = decoder.decode_stream::<T>();

        Ok(std::iter::from_fn(move || {
            if let Err(e) = dbn_stream.advance() {
                return Some(Err(e.into()));
            }

            match dbn_stream.get() {
                Some(rec) => {
                    let record = dbn::RecordRef::from(rec);
                    let instrument_id = match self.resolve_record_instrument_id(
                        &record,
                        instrument_id,
                        &mut metadata_cache,
                    ) {
                        Ok(id) => id,
                        Err(e) => return Some(Err(e)),
                    };

                    let msg = match record.get::<dbn::StatusMsg>() {
                        Some(m) => m,
                        None => return Some(Err(anyhow::anyhow!("Invalid `StatusMsg`"))),
                    };
                    let ts_init = msg.ts_recv.into();

                    match decode_status_msg(msg, instrument_id, Some(ts_init)) {
                        Ok(data) => Some(Ok(data)),
                        Err(e) => Some(Err(e)),
                    }
                }
                None => None,
            }
        }))
    }

    /// Reads imbalance messages from a DBN IMBALANCE schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if reading imbalance records fails.
    pub fn read_imbalance_records<T>(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<DatabentoImbalance>> + '_>
    where
        T: dbn::Record + dbn::HasRType + 'static,
    {
        let decoder = Decoder::from_zstd_file(filepath)?;
        let mut metadata_cache = if instrument_id.is_none() {
            Some(MetadataCache::new(decoder.metadata().clone()))
        } else {
            None
        };
        let mut dbn_stream = decoder.decode_stream::<T>();
        let fixed_instrument_id = instrument_id.is_some();
        let mut fixed_price_precision = price_precision;

        Ok(std::iter::from_fn(move || {
            if let Err(e) = dbn_stream.advance() {
                return Some(Err(e.into()));
            }

            match dbn_stream.get() {
                Some(rec) => {
                    let record = dbn::RecordRef::from(rec);
                    let instrument_id = match self.resolve_record_instrument_id(
                        &record,
                        instrument_id,
                        &mut metadata_cache,
                    ) {
                        Ok(id) => id,
                        Err(e) => return Some(Err(e)),
                    };
                    let resolved_precision = match self.resolve_stream_price_precision(
                        &instrument_id,
                        fixed_instrument_id,
                        &mut fixed_price_precision,
                    ) {
                        Ok(p) => p,
                        Err(e) => return Some(Err(e)),
                    };

                    let msg = match record.get::<dbn::ImbalanceMsg>() {
                        Some(m) => m,
                        None => return Some(Err(anyhow::anyhow!("Invalid `ImbalanceMsg`"))),
                    };
                    let ts_init = msg.ts_recv.into();

                    match decode_imbalance_msg(
                        msg,
                        instrument_id,
                        resolved_precision,
                        Some(ts_init),
                    ) {
                        Ok(data) => Some(Ok(data)),
                        Err(e) => Some(Err(e)),
                    }
                }
                None => None,
            }
        }))
    }

    /// Reads statistics messages from a DBN STATISTICS schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if reading statistics records fails.
    pub fn read_statistics_records<T>(
        &self,
        filepath: &Path,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<DatabentoStatistics>> + '_>
    where
        T: dbn::Record + dbn::HasRType + 'static,
    {
        let decoder = Decoder::from_zstd_file(filepath)?;
        let mut metadata_cache = if instrument_id.is_none() {
            Some(MetadataCache::new(decoder.metadata().clone()))
        } else {
            None
        };
        let mut dbn_stream = decoder.decode_stream::<T>();
        let fixed_instrument_id = instrument_id.is_some();
        let mut fixed_price_precision = price_precision;

        // Loop over skipped records so one unsupported stat_type does not terminate
        // the stream; precheck before precision resolution.
        Ok(std::iter::from_fn(move || {
            loop {
                if let Err(e) = dbn_stream.advance() {
                    return Some(Err(e.into()));
                }

                let rec = dbn_stream.get()?;
                let record = dbn::RecordRef::from(rec);
                let msg = match record.get::<dbn::StatMsg>() {
                    Some(m) => m,
                    None => return Some(Err(anyhow::anyhow!("Invalid `StatMsg`"))),
                };

                if !is_supported_stat_type(msg.stat_type) {
                    log::warn!("Skipping unsupported `stat_type` {}", msg.stat_type);
                    continue;
                }

                let instrument_id = match self.resolve_record_instrument_id(
                    &record,
                    instrument_id,
                    &mut metadata_cache,
                ) {
                    Ok(id) => id,
                    Err(e) => return Some(Err(e)),
                };
                let resolved_precision = match self.resolve_stream_price_precision(
                    &instrument_id,
                    fixed_instrument_id,
                    &mut fixed_price_precision,
                ) {
                    Ok(p) => p,
                    Err(e) => return Some(Err(e)),
                };
                let ts_init = msg.ts_recv.into();

                match decode_statistics_msg(msg, instrument_id, resolved_precision, Some(ts_init)) {
                    Ok(Some(data)) => return Some(Ok(data)),
                    Ok(None) => {}
                    Err(e) => return Some(Err(e)),
                }
            }
        }))
    }

    fn resolve_record_instrument_id(
        &self,
        record: &dbn::RecordRef,
        instrument_id: Option<InstrumentId>,
        metadata_cache: &mut Option<MetadataCache>,
    ) -> anyhow::Result<InstrumentId> {
        if let Some(instrument_id) = instrument_id {
            return Ok(instrument_id);
        }

        let metadata_cache = metadata_cache
            .as_mut()
            .ok_or_else(|| anyhow::anyhow!("missing metadata cache for dynamic instrument id"))?;

        decode_nautilus_instrument_id(
            record,
            metadata_cache,
            &self.publisher_venue_map,
            &self.symbol_venue_map,
        )
    }

    fn resolve_stream_price_precision(
        &self,
        instrument_id: &InstrumentId,
        fixed_instrument_id: bool,
        fixed_price_precision: &mut Option<u8>,
    ) -> anyhow::Result<u8> {
        if let Some(precision) = *fixed_price_precision {
            return Ok(precision);
        }

        let precision = self.resolve_price_precision(instrument_id, None)?;
        if fixed_instrument_id {
            *fixed_price_precision = Some(precision);
        }

        Ok(precision)
    }
}

/// Applies default venue-to-dataset mappings for consolidated Databento feeds.
/// GLBX.MDP3 covers CME Globex exchange MICs; OPRA.PILLAR covers OPRA option venues;
/// EQUS.MINI is the consolidated US equities default.
fn apply_default_venue_dataset_mappings(venue_dataset_map: &mut IndexMap<Venue, Dataset>) {
    let glbx = Dataset::from("GLBX.MDP3");

    for venue in [
        Venue::CBCM(),
        Venue::GLBX(),
        Venue::NYUM(),
        Venue::XCBT(),
        Venue::XCEC(),
        Venue::XCME(),
        Venue::XFXS(),
        Venue::XNYM(),
    ] {
        _ = venue_dataset_map.insert(venue, glbx);
    }

    // publishers.json seeds the consolidated EQUS venue with the unreleased EQUS.PLUS,
    // so pin it to EQUS.MINI, the cheapest released US equities feed.
    _ = venue_dataset_map.insert(Venue::from("EQUS"), Dataset::from("EQUS.MINI"));

    let opra = Dataset::from("OPRA.PILLAR");
    for venue_code in [
        "AMXO", "XBOX", "XCBO", "EMLD", "EDGO", "GMNI", "XISX", "MCRY", "XMIO", "ARCO", "OPRA",
        "MPRL", "XNDQ", "XBXO", "C2OX", "XPHL", "BATO", "MXOP", "SPHR",
    ] {
        _ = venue_dataset_map.insert(Venue::from(venue_code), opra);
    }
}

#[cfg(test)]
mod tests {
    use std::{
        ffi::c_char,
        fs::{File, OpenOptions},
        io::Write,
        path::{Path, PathBuf},
        process,
        sync::atomic::{AtomicU64, Ordering},
    };

    use databento::dbn::encode::EncodeRecord;
    use nautilus_model::{
        enums::BookAction,
        types::{Price, Quantity},
    };
    use rstest::{fixture, rstest};
    use ustr::Ustr;

    use super::*;

    fn test_data_path() -> PathBuf {
        Path::new(env!("CARGO_MANIFEST_DIR")).join("test_data")
    }

    fn mbo_record(action: c_char, flags: u8, sequence: u32) -> dbn::MboMsg {
        let ts_event = 1_609_160_400_000_000_000;
        dbn::MboMsg {
            hd: dbn::RecordHeader::new::<dbn::MboMsg>(dbn::rtype::MBO, 1, 1, ts_event),
            order_id: 42,
            price: 4_800_250_000_000,
            size: 2,
            flags: dbn::FlagSet::new(flags),
            channel_id: 1,
            action,
            side: 'A' as c_char,
            ts_recv: ts_event,
            ts_in_delta: 0,
            sequence,
        }
    }

    fn write_mbo_records(records: &[dbn::MboMsg]) -> PathBuf {
        static FILE_ID: AtomicU64 = AtomicU64::new(0);

        let id = FILE_ID.fetch_add(1, Ordering::Relaxed);
        let path = std::env::temp_dir().join(format!(
            "nautilus-databento-mbo-{}-{id}.dbn.zst",
            process::id(),
        ));
        let metadata = dbn::Metadata::builder()
            .dataset("GLBX.MDP3")
            .schema(Some(dbn::Schema::Mbo))
            .start(0)
            .stype_in(Some(dbn::SType::InstrumentId))
            .stype_out(dbn::SType::InstrumentId)
            .build();
        let file = File::create(&path).unwrap();
        let mut encoder = dbn::encode::dbn::Encoder::with_zstd(file, &metadata).unwrap();
        encoder.encode_records(records).unwrap();
        encoder.flush().unwrap();
        drop(encoder);

        path
    }

    #[fixture]
    fn loader() -> DatabentoDataLoader {
        let publishers_filepath = Path::new(env!("CARGO_MANIFEST_DIR")).join("publishers.json");
        let mut loader = DatabentoDataLoader::new(Some(publishers_filepath)).unwrap();
        // ES futures test data uses precision 2 (USD cents)
        loader.set_price_precision(Symbol::from("ESM4"), 2);
        loader
    }

    #[fixture]
    fn loader_without_seed() -> DatabentoDataLoader {
        let publishers_filepath = Path::new(env!("CARGO_MANIFEST_DIR")).join("publishers.json");
        DatabentoDataLoader::new(Some(publishers_filepath)).unwrap()
    }

    // TODO: Improve the below assertions that we've actually read the records we expected

    #[rstest]
    fn test_set_dataset_venue_mapping(mut loader: DatabentoDataLoader) {
        let dataset = Ustr::from("EQUS.PLUS");
        let venue = Venue::from("XNAS");
        loader.set_dataset_for_venue(dataset, venue);

        let result = loader.get_dataset_for_venue(&venue).unwrap();
        assert_eq!(*result, dataset);
    }

    #[rstest]
    fn test_default_venue_dataset_mappings(loader: DatabentoDataLoader) {
        let xcme = Venue::XCME();
        let result = loader.get_dataset_for_venue(&xcme).unwrap();
        assert_eq!(*result, Ustr::from("GLBX.MDP3"));

        let xcbo = Venue::from("XCBO");
        let result = loader.get_dataset_for_venue(&xcbo).unwrap();
        assert_eq!(*result, Ustr::from("OPRA.PILLAR"));

        let equs = Venue::from("EQUS");
        let result = loader.get_dataset_for_venue(&equs).unwrap();
        assert_eq!(*result, Ustr::from("EQUS.MINI"));
    }

    #[rstest]
    #[case(test_data_path().join("test_data.definition.equity.dbn.zst"))]
    fn test_load_instruments(mut loader: DatabentoDataLoader, #[case] path: PathBuf) {
        let instruments = loader.load_instruments(&path, false, false, None).unwrap();

        assert_eq!(instruments.len(), 2);
        // Definition records auto-populate the precision cache
        assert_eq!(
            loader.get_price_precisions().get(&Symbol::from("ESM4")),
            Some(&2)
        );
    }

    #[rstest]
    fn test_load_instruments_populates_price_precisions_cache(
        mut loader_without_seed: DatabentoDataLoader,
    ) {
        let path = test_data_path().join("test_data.definition.equity.dbn.zst");
        assert!(loader_without_seed.get_price_precisions().is_empty());

        let instruments = loader_without_seed
            .load_instruments(&path, false, false, None)
            .unwrap();

        assert_eq!(instruments.len(), 2);
        for instrument in &instruments {
            let symbol = instrument.id().symbol;
            assert_eq!(
                loader_without_seed.get_price_precisions().get(&symbol),
                Some(&instrument.price_precision()),
                "cache missing or mismatched entry for {symbol}",
            );
        }
    }

    #[rstest]
    fn test_read_records_errors_when_precision_unresolvable(
        loader_without_seed: DatabentoDataLoader,
    ) {
        let path = test_data_path().join("test_data.mbo.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let result = loader_without_seed.load_order_book_deltas(&path, Some(instrument_id), None);

        let err = result.expect_err("expected precision-resolution error");
        let err_msg = format!("{err}");
        assert!(
            err_msg.contains("Could not resolve `price_precision`"),
            "unexpected error message: {err_msg}",
        );
        assert!(
            err_msg.contains("ESM4.GLBX"),
            "error should name the instrument: {err_msg}",
        );
    }

    #[rstest]
    fn test_set_price_precision_unblocks_reads(mut loader_without_seed: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.mbo.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        // Without a seeded precision the read errors
        assert!(
            loader_without_seed
                .load_order_book_deltas(&path, Some(instrument_id), None)
                .is_err()
        );

        loader_without_seed.set_price_precision(Symbol::from("ESM4"), 2);

        let deltas = loader_without_seed
            .load_order_book_deltas(&path, Some(instrument_id), None)
            .unwrap();
        assert_eq!(deltas.len(), 2);
    }

    #[rstest]
    fn test_resolve_price_precision_explicit_arg_overrides_cache(
        mut loader_without_seed: DatabentoDataLoader,
    ) {
        let instrument_id = InstrumentId::from("ESM4.GLBX");
        // Seed a deliberately wrong cache value so we can detect which path is taken
        loader_without_seed.set_price_precision(Symbol::from("ESM4"), 9);

        let explicit = loader_without_seed
            .resolve_price_precision(&instrument_id, Some(2))
            .unwrap();
        assert_eq!(explicit, 2);

        let cached = loader_without_seed
            .resolve_price_precision(&instrument_id, None)
            .unwrap();
        assert_eq!(cached, 9);
    }

    #[rstest]
    fn test_resolve_price_precision_cache_miss_errors(loader_without_seed: DatabentoDataLoader) {
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let err = loader_without_seed
            .resolve_price_precision(&instrument_id, None)
            .expect_err("expected cache-miss error");
        assert!(format!("{err}").contains("Could not resolve `price_precision`"));
    }

    #[rstest]
    fn test_load_order_book_deltas(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.mbo.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let deltas = loader
            .load_order_book_deltas(&path, Some(instrument_id), None)
            .unwrap();

        assert_eq!(deltas.len(), 2);
    }

    #[rstest]
    fn test_read_order_book_deltas_streams_without_collecting(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.mbo.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let count = loader
            .read_order_book_deltas(&path, Some(instrument_id), None)
            .unwrap()
            .map(|result| result.map(|_| 1usize))
            .sum::<anyhow::Result<usize>>()
            .unwrap();

        assert_eq!(count, 2);
    }

    #[rstest]
    #[case::standalone(false)]
    #[case::legacy_inline(true)]
    fn test_load_order_book_deltas_preserves_event_boundary(
        loader: DatabentoDataLoader,
        #[case] inline_last: bool,
    ) {
        let instrument_id = InstrumentId::from("ESM4.GLBX");
        let last = dbn::flags::LAST;
        let mut records = vec![mbo_record(
            'C' as c_char,
            if inline_last { last } else { 0 },
            1,
        )];

        if !inline_last {
            records.push(mbo_record('N' as c_char, last, 2));
        }
        let path = write_mbo_records(&records);

        let deltas = loader
            .load_order_book_deltas(&path, Some(instrument_id), Some(2))
            .unwrap();
        std::fs::remove_file(path).unwrap();

        assert_eq!(deltas.len(), 1);
        assert_eq!(deltas[0].instrument_id, instrument_id);
        assert_eq!(deltas[0].action, BookAction::Delete);
        assert_eq!(deltas[0].order.order_id, 42);
        assert_eq!(deltas[0].flags, last);
        assert_eq!(deltas[0].sequence, 1);
    }

    #[rstest]
    fn test_read_order_book_deltas_drains_before_terminal_error(loader: DatabentoDataLoader) {
        let instrument_id = InstrumentId::from("ESM4.GLBX");
        let path = write_mbo_records(&[mbo_record('C' as c_char, 0, 1)]);
        OpenOptions::new()
            .append(true)
            .open(&path)
            .unwrap()
            .write_all(&[0; 8])
            .unwrap();
        let mut deltas = loader
            .read_order_book_deltas(&path, Some(instrument_id), Some(2))
            .unwrap();

        let delta = deltas.next().unwrap().unwrap();
        let error = deltas
            .next()
            .unwrap()
            .expect_err("expected trailing decode error");
        let end = deltas.next();
        std::fs::remove_file(path).unwrap();

        assert_eq!(delta.instrument_id, instrument_id);
        assert_eq!(delta.order.order_id, 42);
        assert_eq!(delta.flags, 0);
        assert!(format!("{error}").contains("Stream advance error"));
        assert!(end.is_none());
    }

    #[rstest]
    fn test_load_order_book_depth10(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.mbp-10.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let depths = loader
            .load_order_book_depth10(&path, Some(instrument_id), None)
            .unwrap();

        assert_eq!(depths.len(), 2);
    }

    #[rstest]
    fn test_load_quotes(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.mbp-1.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let quotes = loader
            .load_quotes(&path, Some(instrument_id), None)
            .unwrap();

        assert_eq!(quotes.len(), 2);
    }

    #[rstest]
    #[case(test_data_path().join("test_data.bbo-1s.dbn.zst"))]
    #[case(test_data_path().join("test_data.bbo-1m.dbn.zst"))]
    fn test_load_bbo_quotes(loader: DatabentoDataLoader, #[case] path: PathBuf) {
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let quotes = loader
            .load_bbo_quotes(&path, Some(instrument_id), None)
            .unwrap();

        assert_eq!(quotes.len(), 4);
    }

    #[rstest]
    fn test_load_cmbp_quotes(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.cmbp-1.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let quotes = loader
            .load_cmbp_quotes(&path, Some(instrument_id), None)
            .unwrap();

        // Verify exact data count
        assert_eq!(quotes.len(), 2);

        // Verify first quote fields
        let first_quote = &quotes[0];
        assert_eq!(first_quote.instrument_id, instrument_id);
        assert_eq!(first_quote.bid_price, Price::from("3720.25"));
        assert_eq!(first_quote.ask_price, Price::from("3720.50"));
        assert_eq!(first_quote.bid_size, Quantity::from(24));
        assert_eq!(first_quote.ask_size, Quantity::from(11));
        assert_eq!(first_quote.ts_event, 1609160400006136329);
        assert_eq!(first_quote.ts_init, 1609160400006136329);
    }

    #[rstest]
    fn test_load_cbbo_quotes(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.cbbo-1s.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let quotes = loader
            .load_cbbo_quotes(&path, Some(instrument_id), None)
            .unwrap();

        // Verify exact data count
        assert_eq!(quotes.len(), 2);

        // Verify first quote fields
        let first_quote = &quotes[0];
        assert_eq!(first_quote.instrument_id, instrument_id);
        assert_eq!(first_quote.bid_price, Price::from("3720.25"));
        assert_eq!(first_quote.ask_price, Price::from("3720.50"));
        assert_eq!(first_quote.bid_size, Quantity::from(24));
        assert_eq!(first_quote.ask_size, Quantity::from(11));
        assert_eq!(first_quote.ts_event, 1609160400006136329);
        assert_eq!(first_quote.ts_init, 1609160400006136329);
    }

    #[rstest]
    fn test_load_tbbo_trades(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.tbbo.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let trades = loader
            .load_tbbo_trades(&path, Some(instrument_id), None)
            .unwrap();

        assert_eq!(trades.len(), 2);
        assert_eq!(trades[0].instrument_id, instrument_id);
        assert_eq!(trades[0].price, Price::from("3720.25"));
        assert_eq!(trades[0].size, Quantity::from("5"));
    }

    #[rstest]
    fn test_load_tcbbo_trades_rejects_cbbo_fixture(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.cbbo-1s.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let result = loader.load_tcbbo_trades(&path, Some(instrument_id), None);

        assert!(result.is_err());
    }

    #[rstest]
    fn test_load_trades(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.trades.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");
        let trades = loader
            .load_trades(&path, Some(instrument_id), None)
            .unwrap();

        assert_eq!(trades.len(), 2);
    }

    #[rstest]
    // #[case(test_data_path().join("test_data.ohlcv-1d.dbn.zst"))]  // TODO: Empty file (0 records)
    #[case(test_data_path().join("test_data.ohlcv-1h.dbn.zst"))]
    #[case(test_data_path().join("test_data.ohlcv-1m.dbn.zst"))]
    #[case(test_data_path().join("test_data.ohlcv-1s.dbn.zst"))]
    fn test_load_bars(loader: DatabentoDataLoader, #[case] path: PathBuf) {
        let instrument_id = InstrumentId::from("ESM4.GLBX");
        let bars = loader
            .load_bars(&path, Some(instrument_id), None, None)
            .unwrap();

        assert_eq!(bars.len(), 2);
    }

    #[rstest]
    #[case(test_data_path().join("test_data.ohlcv-1s.dbn.zst"))]
    fn test_load_bars_timestamp_on_close_true(loader: DatabentoDataLoader, #[case] path: PathBuf) {
        let instrument_id = InstrumentId::from("ESM4.GLBX");
        let bars = loader
            .load_bars(&path, Some(instrument_id), None, Some(true))
            .unwrap();

        assert_eq!(bars.len(), 2);

        // When bars_timestamp_on_close is true, both ts_event and ts_init should be close time
        for bar in &bars {
            assert_eq!(
                bar.ts_event, bar.ts_init,
                "ts_event and ts_init should both be close time when bars_timestamp_on_close=true"
            );
        }
    }

    #[rstest]
    #[case(test_data_path().join("test_data.ohlcv-1s.dbn.zst"))]
    fn test_load_bars_timestamp_on_close_false(loader: DatabentoDataLoader, #[case] path: PathBuf) {
        let instrument_id = InstrumentId::from("ESM4.GLBX");
        let bars = loader
            .load_bars(&path, Some(instrument_id), None, Some(false))
            .unwrap();

        assert_eq!(bars.len(), 2);

        // When bars_timestamp_on_close is false, ts_event is open time and ts_init is close time
        for bar in &bars {
            assert_ne!(
                bar.ts_event, bar.ts_init,
                "ts_event should be open time and ts_init should be close time when bars_timestamp_on_close=false"
            );
            // For 1-second bars, ts_init (close) should be 1 second after ts_event (open)
            assert_eq!(bar.ts_init.as_u64(), bar.ts_event.as_u64() + 1_000_000_000);
        }
    }

    #[rstest]
    #[case(test_data_path().join("test_data.ohlcv-1s.dbn.zst"), 0)]
    #[case(test_data_path().join("test_data.ohlcv-1s.dbn.zst"), 1)]
    fn test_load_bars_timestamp_comparison(
        loader: DatabentoDataLoader,
        #[case] path: PathBuf,
        #[case] bar_index: usize,
    ) {
        const ONE_SECOND_NS: u64 = 1_000_000_000;

        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let bars_close = loader
            .load_bars(&path, Some(instrument_id), None, Some(true))
            .unwrap();

        let bars_open = loader
            .load_bars(&path, Some(instrument_id), None, Some(false))
            .unwrap();

        assert_eq!(bars_close.len(), bars_open.len());
        assert_eq!(bars_close.len(), 2);

        let bar_close = &bars_close[bar_index];
        let bar_open = &bars_open[bar_index];

        // Bars should have the same OHLCV data
        assert_eq!(bar_close.open, bar_open.open);
        assert_eq!(bar_close.high, bar_open.high);
        assert_eq!(bar_close.low, bar_open.low);
        assert_eq!(bar_close.close, bar_open.close);
        assert_eq!(bar_close.volume, bar_open.volume);

        // The close-timestamped bar should have later timestamp than open-timestamped bar
        // For 1-second bars, this should be exactly 1 second difference
        assert!(
            bar_close.ts_event > bar_open.ts_event,
            "Close-timestamped bar should have later timestamp than open-timestamped bar"
        );

        // The difference should be exactly 1 second (1_000_000_000 nanoseconds) for 1s bars
        assert_eq!(
            bar_close.ts_event.as_u64() - bar_open.ts_event.as_u64(),
            ONE_SECOND_NS,
            "Timestamp difference should be exactly 1 second for 1s bars"
        );
    }

    #[rstest]
    fn test_load_status_records(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.status.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let statuses = loader
            .load_status_records::<dbn::StatusMsg>(&path, Some(instrument_id))
            .unwrap()
            .collect::<anyhow::Result<Vec<_>>>()
            .unwrap();

        // Assert total count matches Python test expectations
        assert_eq!(statuses.len(), 4, "Should load exactly 4 status records");

        // Assert first record fields match Python test expectations
        let first = &statuses[0];
        assert_eq!(first.instrument_id, instrument_id);
        assert_eq!(first.ts_event.as_u64(), 1609110000000000000);
        assert_eq!(first.ts_init.as_u64(), 1609113600000000000);
    }

    #[rstest]
    fn test_read_imbalance_records(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.imbalance.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let imbalances = loader
            .read_imbalance_records::<dbn::ImbalanceMsg>(&path, Some(instrument_id), None)
            .unwrap()
            .collect::<anyhow::Result<Vec<_>>>()
            .unwrap();

        // Assert total count
        assert_eq!(
            imbalances.len(),
            2,
            "Should load exactly 2 imbalance records"
        );

        // Assert first record has required fields
        let first = &imbalances[0];
        assert_eq!(first.instrument_id, instrument_id);
        assert!(
            first.ref_price.as_f64() > 0.0,
            "ref_price should be positive"
        );
        assert!(first.ts_event.as_u64() > 0, "ts_event should be set");
        assert!(first.ts_recv.as_u64() > 0, "ts_recv should be set");
        assert!(first.ts_init.as_u64() > 0, "ts_init should be set");
    }

    #[rstest]
    fn test_read_statistics_records(loader: DatabentoDataLoader) {
        let path = test_data_path().join("test_data.statistics.dbn.zst");
        let instrument_id = InstrumentId::from("ESM4.GLBX");

        let statistics = loader
            .read_statistics_records::<dbn::StatMsg>(&path, Some(instrument_id), None)
            .unwrap()
            .collect::<anyhow::Result<Vec<_>>>()
            .unwrap();

        // Assert total count
        assert_eq!(
            statistics.len(),
            2,
            "Should load exactly 2 statistics records"
        );

        // Assert first record has required fields
        let first = &statistics[0];
        assert_eq!(first.instrument_id, instrument_id);
        assert!(first.ts_event.as_u64() > 0, "ts_event should be set");
        assert!(first.ts_recv.as_u64() > 0, "ts_recv should be set");
        assert!(first.ts_init.as_u64() > 0, "ts_init should be set");
        assert!(first.sequence > 0, "sequence should be positive");
    }
}