atelier_data 0.0.15

Data Artifacts and I/O for the atelier-rs engine
//! Orderbook state persistence
//!
//! Supports saving orderbook snapshots in CSV, JSON, and Parquet formats.
//!
//! # Cargo.toml dependencies
//!
//! ```toml
//! [dependencies]
//! serde = { version = "1.0", features = ["derive"] }
//! serde_json = "1.0"
//! rust_decimal = { version = "1.33", features = ["serde"] }
//!
//! # Optional: only needed for parquet support
//! [dependencies.arrow]
//! version = "53"
//! optional = true
//!
//! [dependencies.parquet]
//! version = "53"
//! optional = true
//!
//! [features]
//! default = []
//! parquet = ["dep:arrow", "dep:parquet"]
//! ```
use crate::utils::current_timestamp_ms;
use crate::{
    errors::PersistError,
    exchanges::Exchange,
    orderbooks::{OrderbookDelta, io},
};
use serde::{Deserialize, Serialize};
use std::path::Path;

/// Output format for orderbook persistence
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OutputFormat {
    Csv,
    Json,
    Parquet,
}

impl OutputFormat {
    /// Infer format from file extension
    pub fn from_path(path: &Path) -> Option<Self> {
        match path.extension()?.to_str()? {
            "csv" => Some(Self::Csv),
            "json" => Some(Self::Json),
            "parquet" | "pq" => Some(Self::Parquet),
            _ => None,
        }
    }
}

/// A single price level for serialization
#[derive(Debug, Clone, Serialize)]
pub struct PriceLevelRecord {
    pub timestamp_ms: u64,
    pub symbol: String,
    pub side: String,  // "bid" or "ask"
    pub level: usize,  // 0 = best, 1 = second best, etc.
    pub price: String, // String to preserve precision
    pub size: String,
}

/// Full orderbook snapshot for JSON serialization
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct OrderbookSnapshot {
    pub timestamp_ms: u64,
    pub symbol: String,
    pub update_id: u64,
    pub sequence: u64,
    pub delta_count: u64,
    pub bid_depth: usize,
    pub ask_depth: usize,
    pub mid_price: Option<String>,
    pub spread: Option<String>,
    pub spread_bps: Option<String>,
    pub volume_imbalance: Option<String>,
    pub total_bid_volume: String,
    pub total_ask_volume: String,
    pub bids: Vec<[String; 2]>, // [[price, size], ...]
    pub asks: Vec<[String; 2]>,
}

// -------------------------------------------------------------------------------
// Public API
// -------------------------------------------------------------------------------

/// Save orderbook state to a file
///
/// Format is inferred from file extension, or can be specified explicitly.
///
/// # Example
/// ```no_run
/// use atelier_data::orderbooks::{delta, persist};
/// let eg_ob = delta::OrderbookDelta::new("BASE_QUOTE");
///
/// persist::save_orderbook_state(&eg_ob, "snapshot.csv", None);
/// persist::save_orderbook_state(&eg_ob, "snapshot.json", None);
/// // This next one requires "parquet" feature
/// persist::save_orderbook_state(&eg_ob, "snapshot.parquet", None);
/// ```
pub fn save_orderbook_state(
    ob: &OrderbookDelta,
    path: impl AsRef<Path>,
    format: Option<OutputFormat>,
) -> Result<(), PersistError> {
    let path = path.as_ref();
    let format = format
        .or_else(|| OutputFormat::from_path(path))
        .ok_or_else(|| {
            PersistError::UnsupportedFormat(
                path.extension()
                    .and_then(|e| e.to_str())
                    .unwrap_or("unknown")
                    .to_string(),
            )
        })?;

    match format {
        OutputFormat::Csv => io::write_csv(ob, path),
        OutputFormat::Json => io::write_json(ob, path),
        OutputFormat::Parquet => io::write_ob_delta_parquet(ob, path),
    }
}

/// Save orderbook state with a timestamp-based filename
///
/// Returns the path to the created file.
pub fn save_orderbook_timestamped(
    ob: &OrderbookDelta,
    ex: &Exchange,
    dir: impl AsRef<Path>,
    format: OutputFormat,
) -> Result<std::path::PathBuf, PersistError> {
    let timestamp = current_timestamp_ms();
    let ext = match format {
        OutputFormat::Csv => "csv",
        OutputFormat::Json => "json",
        OutputFormat::Parquet => "parquet",
    };

    let exchange_name: String = match ex {
        Exchange::Bybit => "bybit".to_string(),
        Exchange::Coinbase => "coinbase".to_string(),
        Exchange::Kraken => "kraken".to_string(),
        Exchange::Binance => "binance".to_string(),
    };

    let filename = format!(
        "{}_{}_{:?}.{}",
        exchange_name.to_lowercase(),
        ob.symbol().to_lowercase(),
        timestamp,
        ext
    );

    let path = dir.as_ref().join(filename);
    save_orderbook_state(ob, &path, Some(format))?;
    Ok(path)
}