//! MBP-10 Data Loader Helper
//!
//! Simple integration layer between DBN MBP-10 parser and OFI feature extraction.
//! Provides synchronous loading wrapper and snapshot window selection for real-time
//! Order Flow Imbalance (OFI) calculation.
//!
//! # Purpose
//!
//! - Load MBP-10 snapshots from DBN files
//! - Provide snapshot windows for OFI calculation
//! - Abstract async DBN parser complexity
//!
//! # Example
//!
//! ```no_run
//! use ml::features::mbp10_loader::load_mbp10_snapshots_sync;
//! use ml::features::ofi_calculator::OFICalculator;
//! use std::path::Path;
//!
//! let snapshots = load_mbp10_snapshots_sync(Path::new("test_data/ES.FUT.mbp10.dbn"))?;
//! let mut calculator = OFICalculator::new();
//!
//! for snapshot in &snapshots {
//!     let features = calculator.calculate(snapshot)?;
//!     // Use features for model input
//! }
//! ```

use data::providers::databento::{dbn_parser::DbnParser, mbp10::{BidAskPair, Mbp10Snapshot}};
use std::path::Path;

use chrono::{DateTime, Utc};

use crate::alternative_bars::ImbalanceBarSampler;
use crate::ofi_calculator::OFICalculator;
use crate::trades_loader::load_trades_sync;
use crate::OHLCVBar;
use crate::MLError;

/// Load MBP-10 snapshots from DBN file (synchronous wrapper)
///
/// This is a blocking wrapper around the async `DbnParser::parse_mbp10_file()`.
/// Use this for simple synchronous contexts like testing or batch processing.
///
/// # Arguments
///
/// * `file_path` - Path to the DBN file containing MBP-10 data
///
/// # Returns
///
/// * `Ok(Vec<Mbp10Snapshot>)` - All snapshots from the file
/// * `Err(MLError)` - If file cannot be read or parsed
///
/// # Performance
///
/// - Target: <10ms for typical 180-day files
/// - Memory: ~1KB per snapshot (10 levels × 48 bytes)
///
pub fn load_mbp10_snapshots_sync(file_path: &Path) -> Result<Vec<Mbp10Snapshot>, MLError> {
    let parser = DbnParser::new().map_err(|e| {
        MLError::InsufficientData(format!("Failed to create DBN parser: {}", e))
    })?;

    // Use tokio runtime for async operation
    let runtime = tokio::runtime::Runtime::new().map_err(|e| {
        MLError::InsufficientData(format!("Failed to create async runtime: {}", e))
    })?;

    runtime.block_on(async {
        parser
            .parse_mbp10_file(file_path)
            .await
            .map_err(|e| MLError::InsufficientData(format!("Failed to parse MBP-10 file: {}", e)))
    })
}

/// Compute OFI features from an MBP-10 file in a single streaming pass.
///
/// This is ~10x faster than `load_mbp10_snapshots_sync` + separate OFI iteration
/// because it avoids materializing millions of intermediate `Mbp10Snapshot` clones.
/// The parser feeds snapshots by reference directly to the OFI calculator.
///
/// # Arguments
///
/// * `file_path` - Path to the .dbn or .dbn.zst MBP-10 file
///
/// # Returns
///
/// * `Ok(Vec<[f64; 8]>)` - OFI feature arrays (one per aggregated snapshot)
/// * `Err(MLError)` - If file cannot be read or parsed
///
pub fn compute_ofi_from_file(file_path: &Path) -> Result<Vec<[f64; 8]>, MLError> {
    let parser = DbnParser::new().map_err(|e| {
        MLError::InsufficientData(format!("Failed to create DBN parser: {}", e))
    })?;

    let mut calculator = OFICalculator::new();
    let mut features = Vec::new();

    let snapshot_count = parser
        .parse_mbp10_streaming(file_path, 100, |snapshot| {
            if let Ok(f) = calculator.calculate(snapshot) {
                features.push(f.to_array());
            } // Skip failed calculations (e.g. first snapshot with no prev)
        })
        .map_err(|e| MLError::InsufficientData(format!("MBP-10 streaming parse failed: {}", e)))?;

    tracing::debug!(
        "Computed {} OFI features from {} snapshots in {:?}",
        features.len(),
        snapshot_count,
        file_path.file_name().unwrap_or_default()
    );

    Ok(features)
}

/// Compute OFI features from MBP-10 file enriched with trade data.
///
/// Interleaves trade events (from a `.dbn.zst` trades file) into the MBP-10
/// streaming pipeline. Before each MBP-10 snapshot is processed, all trades
/// with timestamps up to that snapshot are fed into the OFI calculator,
/// populating VPIN, Kyle's Lambda, and trade imbalance with real trade data
/// instead of leaving them at zero.
///
/// Falls back to `compute_ofi_from_file()` if the trade file doesn't exist
/// or fails to load.
pub fn compute_ofi_with_trades(
    mbp10_file: &Path,
    trades_file: &Path,
) -> Result<Vec<[f64; 8]>, MLError> {
    // Load all trades upfront (sorted by timestamp). Trade files are much
    // smaller than MBP-10 (~120 MB vs ~5+ GB), so this is fine in memory.
    let trades = match load_trades_sync(trades_file) {
        Ok(t) => {
            tracing::info!(
                "Loaded {} trades from {} for OFI enrichment",
                t.len(),
                trades_file.file_name().unwrap_or_default().to_string_lossy()
            );
            t
        }
        Err(e) => {
            tracing::warn!(
                "Failed to load trades from {}: {} \u{2014} falling back to MBP-10 only",
                trades_file.display(),
                e
            );
            return compute_ofi_from_file(mbp10_file);
        }
    };

    let parser = DbnParser::new().map_err(|e| {
        MLError::InsufficientData(format!("Failed to create DBN parser: {}", e))
    })?;

    let mut calculator = OFICalculator::new();
    let mut features = Vec::new();
    let mut trade_cursor: usize = 0;

    let snapshot_count = parser
        .parse_mbp10_streaming(mbp10_file, 100, |snapshot| {
            // Feed all trades with timestamp <= this snapshot's timestamp.
            // This ensures VPIN/Kyle's Lambda/trade_imbalance are populated
            // before the OFI calculator processes the snapshot.
            while trade_cursor < trades.len()
                && trades[trade_cursor].timestamp <= snapshot.timestamp
            {
                let t = &trades[trade_cursor];
                calculator.feed_trade(t.price, t.volume, t.is_buy);
                trade_cursor += 1;
            }

            if let Ok(f) = calculator.calculate(snapshot) {
                features.push(f.to_array());
            } // Skip failed calculations (e.g. first snapshot)
        })
        .map_err(|e| MLError::InsufficientData(format!("MBP-10 streaming parse failed: {}", e)))?;

    tracing::info!(
        "Computed {} OFI features from {} snapshots + {} trades (cursor at {}/{})",
        features.len(),
        snapshot_count,
        trades.len(),
        trade_cursor,
        trades.len()
    );

    Ok(features)
}

/// Find matching trade file for an MBP-10 file in the trades directory.
///
/// Matches by filename stem: `ES.FUT_2024-Q1.dbn.zst` in MBP-10 dir
/// → looks for `ES.FUT_2024-Q1.dbn.zst` in trades dir (maintaining
/// the same subdirectory structure).
fn find_matching_trade_file(mbp10_file: &Path, trades_dir: &Path) -> Option<std::path::PathBuf> {
    let file_name = mbp10_file.file_name()?;

    // Try exact match: trades_dir/SYMBOL/filename
    // MBP-10 files are typically under mbp10_dir/ES.FUT/ES.FUT_2024-Q1.dbn.zst
    // Trades would be under trades_dir/ES.FUT/ES.FUT_2024-Q1.dbn.zst
    if let Some(parent) = mbp10_file.parent() {
        if let Some(symbol_dir) = parent.file_name() {
            let candidate = trades_dir.join(symbol_dir).join(file_name);
            if candidate.exists() {
                return Some(candidate);
            }
        }
    }

    // Fallback: flat structure — trades_dir/filename
    let candidate = trades_dir.join(file_name);
    if candidate.exists() {
        return Some(candidate);
    }

    None
}

/// Load OFI features from all MBP-10 files in a directory using parallel streaming.
///
/// Files are processed concurrently with rayon. Each file streams through the DBN parser
/// and computes OFI features inline — no intermediate `Vec<Mbp10Snapshot>` allocation.
///
/// # Arguments
///
/// * `mbp10_dir` - Directory containing .dbn or .dbn.zst MBP-10 files
/// * `file_collector` - Function to collect files from the directory (supports both
///   flat and recursive collection patterns used by DQN/PPO adapters)
///
/// # Returns
///
/// * `Some(Vec<[f64; 8]>)` - Concatenated OFI features from all files (file order preserved)
/// * `None` - If directory doesn't exist, has no files, or all files failed
///
#[allow(clippy::cognitive_complexity)]
pub fn load_ofi_features_parallel<F>(
    mbp10_dir: &Path,
    trades_dir: Option<&Path>,
    file_collector: F,
) -> Option<Vec<[f64; 8]>>
where
    F: FnOnce(&Path) -> Vec<std::path::PathBuf>,
{
    use rayon::prelude::*;

    if !mbp10_dir.exists() {
        tracing::info!(
            "No MBP10 directory at {}, OFI features will be zero-padded",
            mbp10_dir.display()
        );
        return None;
    }

    let mbp10_files = file_collector(mbp10_dir);
    if mbp10_files.is_empty() {
        tracing::info!("No MBP10 .dbn files found in {}", mbp10_dir.display());
        return None;
    }

    tracing::info!(
        "Loading OFI features from {} MBP10 files (parallel streaming{})",
        mbp10_files.len(),
        if trades_dir.is_some() { " + trades" } else { "" }
    );

    let start = std::time::Instant::now();

    // Process files in parallel — each file gets its own OFI calculator (independent state).
    // If a trades directory is provided, each MBP-10 file looks for a matching trade file
    // to enrich VPIN, Kyle's Lambda, and trade imbalance with real trade data.
    let per_file_results: Vec<_> = mbp10_files
        .par_iter()
        .map(|file| {
            let result = if let Some(t_dir) = trades_dir {
                if let Some(trade_file) = find_matching_trade_file(file, t_dir) {
                    compute_ofi_with_trades(file, &trade_file)
                } else {
                    tracing::debug!(
                        "No matching trade file for {}, using MBP-10 only",
                        file.file_name().unwrap_or_default().to_string_lossy()
                    );
                    compute_ofi_from_file(file)
                }
            } else {
                compute_ofi_from_file(file)
            };

            match result {
                Ok(features) => {
                    tracing::info!(
                        "  {} \u{2192} {} OFI features",
                        file.file_name().unwrap_or_default().to_string_lossy(),
                        features.len()
                    );
                    features
                }
                Err(e) => {
                    tracing::warn!("Failed to compute OFI from {}: {}", file.display(), e);
                    Vec::new()
                }
            }
        })
        .collect();

    // Concatenate results in file order
    let total_features: usize = per_file_results.iter().map(|v| v.len()).sum();
    let mut all_ofi = Vec::with_capacity(total_features);
    for features in per_file_results {
        all_ofi.extend_from_slice(&features);
    }

    if all_ofi.is_empty() {
        return None;
    }

    let elapsed = start.elapsed();
    tracing::info!(
        "Computed {} OFI feature vectors from {} files in {:.1}s ({:.0} features/sec)",
        all_ofi.len(),
        mbp10_files.len(),
        elapsed.as_secs_f64(),
        all_ofi.len() as f64 / elapsed.as_secs_f64()
    );

    Some(all_ofi)
}

/// Trade tick extracted from MBP-10 data.
///
/// MBP-10 records include order book updates AND trades. When an MBP-10 record
/// has `action == b'T'` it represents a trade execution. We extract the price,
/// size, timestamp, and buy/sell classification for feeding into the
/// `ImbalanceBarSampler`.
#[derive(Debug, Clone)]
pub struct Mbp10Trade {
    /// Trade price (f64, converted from fixed-point)
    pub price: f64,
    /// Trade volume in contracts
    pub volume: f64,
    /// Event timestamp
    pub timestamp: DateTime<Utc>,
    /// True if buyer-initiated (side == 'B'), false otherwise
    pub is_buy: bool,
}

/// Extract trade ticks from MBP-10 snapshots using consecutive snapshot diffs.
///
/// MBP-10 snapshots don't directly carry trade action/side info (those live on
/// the raw `dbn::Mbp10Msg` records). This function uses a heuristic:
/// when `trade_count` increases between consecutive snapshots, a trade occurred.
/// Direction is inferred by comparing mid-prices (tick rule).
///
/// For higher-fidelity trade extraction, use [`extract_trades_from_dbn_file`]
/// which reads raw MBP-10 records with their action/side fields.
pub fn extract_trades_from_snapshots(snapshots: &[Mbp10Snapshot]) -> Vec<Mbp10Trade> {
    let mut trades = Vec::new();
    let mut prev_trade_count: u32 = 0;
    let mut prev_mid: f64 = 0.0;

    for snap in snapshots {
        let mid = snap.mid_price();
        if mid <= 0.0 {
            continue;
        }

        let new_trades = snap.trade_count.saturating_sub(prev_trade_count);
        if new_trades > 0 && prev_mid > 0.0 {
            // Tick rule: price up => buy, price down => sell, unchanged => previous direction
            let is_buy = mid >= prev_mid;
            // Best ask size as proxy for trade volume (actual volume not in snapshot)
            let volume = if !snap.levels.is_empty() {
                snap.levels[0].bid_sz.max(1) as f64
            } else {
                1.0
            };

            let ts_secs = (snap.timestamp / 1_000_000_000) as i64;
            let ts_nanos = (snap.timestamp % 1_000_000_000) as u32;
            if let Some(dt) = DateTime::<Utc>::from_timestamp(ts_secs, ts_nanos) {
                trades.push(Mbp10Trade {
                    price: mid,
                    volume,
                    timestamp: dt,
                    is_buy,
                });
            }
        }

        prev_trade_count = snap.trade_count;
        prev_mid = mid;
    }

    trades
}

/// Extract trade ticks directly from a DBN MBP-10 file.
///
/// Reads raw `Mbp10Msg` records and filters for trade actions (`action == b'T'`).
/// This provides higher-fidelity trades than [`extract_trades_from_snapshots`]
/// because it uses the actual action and side fields from each record.
pub fn extract_trades_from_dbn_file(file_path: &Path) -> Result<Vec<Mbp10Trade>, MLError> {
    use dbn::decode::{DbnDecoder, DecodeRecordRef};
    use dbn::RecordRefEnum;
    use std::fs::File;
    use std::io::BufReader;

    let file = File::open(file_path).map_err(|e| {
        MLError::InsufficientData(format!("Failed to open MBP-10 file {:?}: {}", file_path, e))
    })?;

    let is_zstd = file_path.to_string_lossy().ends_with(".dbn.zst");
    let reader: Box<dyn std::io::Read> = if is_zstd {
        Box::new(
            zstd::Decoder::new(BufReader::new(file)).map_err(|e| {
                MLError::InsufficientData(format!("Failed to create zstd decoder: {}", e))
            })?,
        )
    } else {
        Box::new(BufReader::new(file))
    };

    let mut decoder = DbnDecoder::new(reader).map_err(|e| {
        MLError::InsufficientData(format!("Failed to create DBN decoder: {}", e))
    })?;

    let mut trades = Vec::new();
    let mut last_mid: f64 = 0.0;

    loop {
        match decoder.decode_record_ref() {
            Ok(Some(record)) => {
                let record_enum = record.as_enum().map_err(|e| {
                    MLError::InsufficientData(format!("Failed to convert record: {}", e))
                })?;

                if let RecordRefEnum::Mbp10(mbp10) = record_enum {
                    let action = mbp10.action as u8;

                    // Compute mid-price from this record's top-of-book levels
                    // MBP-10 records carry 10 levels; index 0 is best bid/ask
                    let bid_px = BidAskPair::price_to_f64(mbp10.levels[0].bid_px);
                    let ask_px = BidAskPair::price_to_f64(mbp10.levels[0].ask_px);
                    let mid = if bid_px > 0.0 && ask_px > 0.0 {
                        (bid_px + ask_px) / 2.0
                    } else {
                        last_mid
                    };
                    if mid > 0.0 {
                        last_mid = mid;
                    }

                    // Trade action: 'T' (0x54) = trade execution
                    if action == b'T' {
                        let price = mbp10.price as f64 * 1e-9;
                        let volume = mbp10.size as f64;
                        let is_buy = mbp10.side == b'B' as i8;

                        let ts_secs = (mbp10.hd.ts_event / 1_000_000_000) as i64;
                        let ts_nanos = (mbp10.hd.ts_event % 1_000_000_000) as u32;
                        if let Some(dt) = DateTime::<Utc>::from_timestamp(ts_secs, ts_nanos) {
                            if price > 0.0 && volume > 0.0 {
                                trades.push(Mbp10Trade {
                                    price,
                                    volume,
                                    timestamp: dt,
                                    is_buy,
                                });
                            }
                        }
                    }
                }
            }
            Ok(None) => break,
            Err(e) => {
                return Err(MLError::InsufficientData(format!(
                    "DBN decode error: {}",
                    e
                )));
            }
        }
    }

    Ok(trades)
}

/// Recursively collect all .dbn and .dbn.zst files from a directory.
fn collect_dbn_files(dir: &Path) -> Vec<std::path::PathBuf> {
    let mut files = Vec::new();
    if let Ok(entries) = std::fs::read_dir(dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                files.extend(collect_dbn_files(&path));
            } else if path.extension().and_then(|s| s.to_str()) == Some("dbn")
                || path.to_string_lossy().ends_with(".dbn.zst")
            {
                files.push(path);
            }
        }
    }
    files
}

// ══════════════════════════════════════════════════════════════════════════════
// Imbalance Bar Cache
//
// Computing imbalance bars from MBP-10 data requires decompressing ~19GB of
// .dbn.zst files (~450s). This cache stores the resulting `Vec<OHLCVBar>` to
// `/tmp/.foxhunt_imbalance_cache/<hash>.bars.bin` using bincode.
//
// Cache key: DefaultHasher of (mbp10_dir, symbol, threshold bits, alpha bits).
// Invalidation: cache mtime vs. newest .dbn/.dbn.zst file in the source dir.
// ══════════════════════════════════════════════════════════════════════════════

/// Compute a cache path for an imbalance bar set.
fn imbalance_cache_path(mbp10_dir: &Path, symbol: &str, threshold: f64, alpha: f64) -> std::path::PathBuf {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    let mut hasher = DefaultHasher::new();
    mbp10_dir.hash(&mut hasher);
    symbol.hash(&mut hasher);
    threshold.to_bits().hash(&mut hasher);
    alpha.to_bits().hash(&mut hasher);
    let hash = hasher.finish();
    std::path::PathBuf::from(format!("/tmp/.foxhunt_imbalance_cache/{:016x}.bars.bin", hash))
}

/// Check whether the cache file is newer than every `.dbn`/`.dbn.zst` file under `source_dir`.
fn is_imbalance_cache_valid(cache: &Path, source_dir: &Path) -> bool {
    let cache_mtime = match std::fs::metadata(cache) {
        Ok(m) => match m.modified() {
            Ok(t) => t,
            Err(_) => return false,
        },
        Err(_) => return false,
    };

    fn walk_dbn_mtimes(dir: &Path, newest: &mut std::time::SystemTime) -> std::io::Result<()> {
        for entry in std::fs::read_dir(dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.is_dir() {
                walk_dbn_mtimes(&path, newest)?;
            } else if path.extension().and_then(|s| s.to_str()) == Some("dbn")
                || path.to_string_lossy().ends_with(".dbn.zst")
            {
                if let Ok(meta) = std::fs::metadata(&path) {
                    if let Ok(mtime) = meta.modified() {
                        if mtime > *newest {
                            *newest = mtime;
                        }
                    }
                }
            }
        }
        Ok(())
    }

    let mut newest_source = std::time::SystemTime::UNIX_EPOCH;
    if walk_dbn_mtimes(source_dir, &mut newest_source).is_err() {
        return false;
    }

    // Cache is valid if it was created after all source files were last modified
    cache_mtime >= newest_source
}

/// Load MBP-10 files and convert to imbalance bars via `ImbalanceBarSampler`.
///
/// Scans `mbp10_dir/symbol/` for `.dbn` and `.dbn.zst` files, loads each,
/// extracts trades (action == 'T'), feeds into adaptive `ImbalanceBarSampler`.
///
/// Results are cached to `/tmp/.foxhunt_imbalance_cache/<hash>.bars.bin`.
/// Subsequent calls with the same parameters return in <1s if no source files
/// have changed.
///
/// Returns sorted `OHLCVBar`s or error if directory doesn't exist.
/// No fallback -- fails loudly if the data isn't there.
pub fn mbp10_to_imbalance_bars(
    mbp10_dir: &Path,
    symbol: &str,
    threshold: f64,
    ewma_alpha: f64,
) -> Result<Vec<OHLCVBar>, MLError> {
    let symbol_dir = mbp10_dir.join(symbol);
    if !symbol_dir.exists() {
        return Err(MLError::InsufficientData(format!(
            "MBP-10 directory not found: {}. Set data_source='ohlcv' or provide MBP-10 data.",
            symbol_dir.display()
        )));
    }

    // ── Cache check ──────────────────────────────────────────────────────
    let cache_file = imbalance_cache_path(mbp10_dir, symbol, threshold, ewma_alpha);
    if is_imbalance_cache_valid(&cache_file, &symbol_dir) {
        let load_start = std::time::Instant::now();
        match load_imbalance_cache(&cache_file) {
            Ok(bars) => {
                let elapsed_ms = load_start.elapsed().as_secs_f64() * 1000.0;
                tracing::info!(
                    "Imbalance bar cache HIT: loaded {} bars in {:.1}ms",
                    bars.len(),
                    elapsed_ms
                );
                return Ok(bars);
            }
            Err(e) => {
                tracing::warn!("Imbalance bar cache read failed (recomputing): {}", e);
            }
        }
    }

    // ── Cache miss — compute from scratch ────────────────────────────────
    let mut dbn_files = collect_dbn_files(&symbol_dir);
    dbn_files.sort();

    if dbn_files.is_empty() {
        return Err(MLError::InsufficientData(format!(
            "No .dbn/.dbn.zst files found in {}",
            symbol_dir.display()
        )));
    }

    tracing::info!(
        "Imbalance bar cache MISS: computing from {} MBP-10 files...",
        dbn_files.len()
    );

    // Collect all trades from all files
    let mut all_trades: Vec<Mbp10Trade> = Vec::new();
    for file_path in &dbn_files {
        match extract_trades_from_dbn_file(file_path) {
            Ok(trades) => {
                tracing::debug!(
                    "  {} -> {} trades",
                    file_path.file_name().unwrap_or_default().to_string_lossy(),
                    trades.len()
                );
                all_trades.extend(trades);
            }
            Err(e) => {
                tracing::warn!(
                    "Failed to extract trades from {}: {}",
                    file_path.display(),
                    e
                );
            }
        }
    }

    if all_trades.is_empty() {
        return Err(MLError::InsufficientData(format!(
            "No trade ticks extracted from MBP-10 files in {}. \
             Files may contain only book updates without trade events.",
            symbol_dir.display()
        )));
    }

    // Sort by timestamp (files should be chronological but ensure correctness)
    all_trades.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));

    tracing::info!(
        "Extracted {} trade ticks, feeding into ImbalanceBarSampler (threshold={}, alpha={})",
        all_trades.len(),
        threshold,
        ewma_alpha,
    );

    // Initialize sampler with first trade
    let first = &all_trades[0];
    let mut sampler =
        ImbalanceBarSampler::new_with_ewma(first.price, threshold, first.timestamp, ewma_alpha);

    let mut bars: Vec<OHLCVBar> = Vec::new();
    for trade in &all_trades {
        if let Some(bar) = sampler.update(trade.price, trade.volume, trade.timestamp) {
            bars.push(bar);
        }
    }

    // Sort by timestamp (should already be ordered)
    bars.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));

    tracing::info!(
        "ImbalanceBarSampler produced {} bars from {} trade ticks (ratio 1:{:.0})",
        bars.len(),
        all_trades.len(),
        if bars.is_empty() {
            0.0
        } else {
            all_trades.len() as f64 / bars.len() as f64
        },
    );

    // ── Save to cache ────────────────────────────────────────────────────
    if let Err(e) = save_imbalance_cache(&cache_file, &bars) {
        tracing::warn!("Failed to write imbalance bar cache (non-fatal): {}", e);
    } else {
        tracing::info!("Cached {} bars to {}", bars.len(), cache_file.display());
    }

    Ok(bars)
}

/// Deserialize cached imbalance bars from a bincode file.
fn load_imbalance_cache(path: &Path) -> Result<Vec<OHLCVBar>, MLError> {
    use std::io::BufReader;
    let file = std::fs::File::open(path).map_err(|e| {
        MLError::InsufficientData(format!("Failed to open imbalance cache {:?}: {}", path, e))
    })?;
    let reader = BufReader::new(file);
    bincode::deserialize_from(reader).map_err(|e| {
        MLError::InsufficientData(format!(
            "Failed to deserialize imbalance cache {:?}: {}",
            path, e
        ))
    })
}

/// Serialize imbalance bars to a bincode file, creating the cache directory if needed.
fn save_imbalance_cache(path: &Path, bars: &[OHLCVBar]) -> Result<(), MLError> {
    use std::io::BufWriter;
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).map_err(|e| {
            MLError::InsufficientData(format!(
                "Failed to create imbalance cache dir {:?}: {}",
                parent, e
            ))
        })?;
    }
    let file = std::fs::File::create(path).map_err(|e| {
        MLError::InsufficientData(format!(
            "Failed to create imbalance cache {:?}: {}",
            path, e
        ))
    })?;
    let writer = BufWriter::new(file);
    bincode::serialize_into(writer, bars).map_err(|e| {
        MLError::InsufficientData(format!(
            "Failed to serialize imbalance cache {:?}: {}",
            path, e
        ))
    })
}

/// Get snapshots for a specific timestamp window
///
/// Returns a slice of snapshots starting from the first snapshot at or before
/// the target timestamp. Used for calculating OFI features that require
/// forward-looking context (e.g., next N snapshots after a given timestamp).
///
/// # Arguments
///
/// * `snapshots` - All available snapshots (must be sorted by timestamp)
/// * `target_ts` - Target timestamp (nanoseconds since Unix epoch)
/// * `window_size` - Number of snapshots to include in window
///
/// # Returns
///
/// Slice of up to `window_size` snapshots starting from the first snapshot
/// with timestamp >= `target_ts`. Returns empty slice if no snapshots exist
/// at or after the target timestamp.
///
/// # Algorithm
///
/// - Simple linear search over the sorted snapshot slice.
/// - Returns snapshots starting from the first one with timestamp >= `target_ts`.
///
/// # Example
///
/// ```ignore
/// // Snapshots at times: [1000, 2000, 3000, 4000]
/// // target_ts = 2000, window_size = 2
/// // Returns: [2000, 3000]
/// ```
///
pub fn get_snapshots_for_timestamp(
    snapshots: &[Mbp10Snapshot],
    target_ts: u64,
    window_size: usize,
) -> &[Mbp10Snapshot] {
    if snapshots.is_empty() {
        return &[];
    }

    // Binary search for the first snapshot with timestamp >= target_ts
    // Snapshots are pre-sorted by timestamp (from DBN file ordering)
    let idx = snapshots.partition_point(|s| s.timestamp < target_ts);
    if idx >= snapshots.len() {
        return &[]; // Target timestamp is after all snapshots
    }
    let end_idx = (idx + window_size).min(snapshots.len());
    &snapshots[idx..end_idx]
}

/// Get the most recent N snapshots ending at the given index
///
/// Useful for calculating features that require a rolling window of recent data.
///
/// # Arguments
///
/// * `snapshots` - All available snapshots
/// * `end_idx` - Index of the last snapshot to include (exclusive)
/// * `window_size` - Number of snapshots to include
///
/// # Returns
///
/// Slice of up to `window_size` snapshots ending at `end_idx`
///
pub fn get_recent_snapshots(
    snapshots: &[Mbp10Snapshot],
    end_idx: usize,
    window_size: usize,
) -> &[Mbp10Snapshot] {
    if snapshots.is_empty() || end_idx == 0 {
        return &[];
    }

    let actual_end = end_idx.min(snapshots.len());
    let start_idx = actual_end.saturating_sub(window_size);

    &snapshots[start_idx..actual_end]
}

#[cfg(test)]
mod tests {
    use super::*;
    use data::providers::databento::mbp10::BidAskPair;

    fn create_test_snapshot(timestamp: u64, bid_px: i64) -> Mbp10Snapshot {
        let levels = vec![BidAskPair {
            bid_px,
            bid_sz: 100,
            bid_ct: 5,
            ask_px: bid_px + 1000000000, // 1 tick higher
            ask_sz: 120,
            ask_ct: 6,
        }];

        Mbp10Snapshot::new("ES.FUT".to_owned(), timestamp, levels, 0, 0)
    }

    #[test]
    fn test_get_snapshots_for_timestamp_exact_match() {
        let snapshots = vec![
            create_test_snapshot(1000, 150000000000000),
            create_test_snapshot(2000, 150010000000000),
            create_test_snapshot(3000, 150020000000000),
            create_test_snapshot(4000, 150030000000000),
        ];

        let window = get_snapshots_for_timestamp(&snapshots, 2000, 2);
        assert_eq!(window.len(), 2);
        assert_eq!(window[0].timestamp, 2000);
        assert_eq!(window[1].timestamp, 3000);
    }

    #[test]
    fn test_get_snapshots_for_timestamp_between() {
        let snapshots = vec![
            create_test_snapshot(1000, 150000000000000),
            create_test_snapshot(2000, 150010000000000),
            create_test_snapshot(3000, 150020000000000),
        ];

        // Target timestamp between 2000 and 3000
        let window = get_snapshots_for_timestamp(&snapshots, 2500, 2);
        assert_eq!(window.len(), 1); // Only snapshot at 3000 remains
        assert_eq!(window[0].timestamp, 3000);
    }

    #[test]
    fn test_get_snapshots_for_timestamp_empty() {
        let snapshots: Vec<Mbp10Snapshot> = vec![];
        let window = get_snapshots_for_timestamp(&snapshots, 2000, 2);
        assert_eq!(window.len(), 0);
    }

    #[test]
    fn test_get_snapshots_for_timestamp_before_all() {
        let snapshots = vec![
            create_test_snapshot(2000, 150000000000000),
            create_test_snapshot(3000, 150010000000000),
        ];

        // Target is before all snapshots, should return first 2 snapshots
        let window = get_snapshots_for_timestamp(&snapshots, 1000, 2);
        assert_eq!(window.len(), 2);
        assert_eq!(window[0].timestamp, 2000);
        assert_eq!(window[1].timestamp, 3000);
    }

    #[test]
    fn test_get_snapshots_for_timestamp_after_all() {
        let snapshots = vec![
            create_test_snapshot(1000, 150000000000000),
            create_test_snapshot(2000, 150010000000000),
        ];

        let window = get_snapshots_for_timestamp(&snapshots, 5000, 2);
        assert_eq!(window.len(), 0); // No snapshots after target
    }

    #[test]
    fn test_get_snapshots_window_clipping() {
        let snapshots = vec![
            create_test_snapshot(1000, 150000000000000),
            create_test_snapshot(2000, 150010000000000),
            create_test_snapshot(3000, 150020000000000),
        ];

        // Request 10 snapshots but only 2 available at/after target (2000, 3000)
        let window = get_snapshots_for_timestamp(&snapshots, 2000, 10);
        assert_eq!(window.len(), 2); // 2 snapshots at/after 2000: [2000, 3000]
        assert_eq!(window[0].timestamp, 2000);
        assert_eq!(window[1].timestamp, 3000);
    }

    #[test]
    fn test_get_recent_snapshots_full_window() {
        let snapshots = vec![
            create_test_snapshot(1000, 150000000000000),
            create_test_snapshot(2000, 150010000000000),
            create_test_snapshot(3000, 150020000000000),
            create_test_snapshot(4000, 150030000000000),
        ];

        let window = get_recent_snapshots(&snapshots, 3, 2);
        assert_eq!(window.len(), 2);
        assert_eq!(window[0].timestamp, 2000);
        assert_eq!(window[1].timestamp, 3000);
    }

    #[test]
    fn test_get_recent_snapshots_partial_window() {
        let snapshots = vec![
            create_test_snapshot(1000, 150000000000000),
            create_test_snapshot(2000, 150010000000000),
        ];

        // Request 5 snapshots but only 2 available
        let window = get_recent_snapshots(&snapshots, 2, 5);
        assert_eq!(window.len(), 2);
        assert_eq!(window[0].timestamp, 1000);
        assert_eq!(window[1].timestamp, 2000);
    }

    #[test]
    fn test_get_recent_snapshots_empty() {
        let snapshots: Vec<Mbp10Snapshot> = vec![];
        let window = get_recent_snapshots(&snapshots, 0, 2);
        assert_eq!(window.len(), 0);
    }

    #[test]
    fn test_get_recent_snapshots_zero_end_idx() {
        let snapshots = vec![
            create_test_snapshot(1000, 150000000000000),
            create_test_snapshot(2000, 150010000000000),
        ];

        let window = get_recent_snapshots(&snapshots, 0, 2);
        assert_eq!(window.len(), 0);
    }

    #[test]
    fn test_get_recent_snapshots_one_element() {
        let snapshots = vec![create_test_snapshot(1000, 150000000000000)];

        let window = get_recent_snapshots(&snapshots, 1, 1);
        assert_eq!(window.len(), 1);
        assert_eq!(window[0].timestamp, 1000);
    }

    #[test]
    fn test_extract_trades_from_snapshots() {
        // Create snapshots with increasing trade_count and rising prices
        let snapshots = vec![
            {
                let mut s = create_test_snapshot(1_000_000_000, 150000000000000); // 150.0
                s.trade_count = 0;
                s
            },
            {
                let mut s = create_test_snapshot(2_000_000_000, 150010000000000); // 150.01
                s.trade_count = 1;
                s
            },
            {
                let mut s = create_test_snapshot(3_000_000_000, 150020000000000); // 150.02
                s.trade_count = 3;
                s
            },
            {
                let mut s = create_test_snapshot(4_000_000_000, 149990000000000); // 149.99
                s.trade_count = 4;
                s
            },
        ];

        let trades = extract_trades_from_snapshots(&snapshots);
        // 3 transitions where trade_count increases (snap 0->1, 1->2, 2->3)
        assert_eq!(trades.len(), 3);
        // First trade: price went up => is_buy=true
        assert!(trades[0].is_buy);
        // Third trade: price went down => is_buy=false
        assert!(!trades[2].is_buy);
    }

    #[test]
    fn test_mbp10_to_imbalance_bars_missing_dir() {
        let result = mbp10_to_imbalance_bars(
            Path::new("/nonexistent/path"),
            "ES.FUT",
            100.0,
            0.1,
        );
        assert!(result.is_err());
        let err_msg = format!("{}", result.unwrap_err());
        assert!(err_msg.contains("MBP-10 directory not found"), "Error: {err_msg}");
    }

    /// Integration test: loads real MBP-10 data and produces imbalance bars.
    /// Reads threshold/alpha from config/training/dqn-smoketest.toml (single source of truth).
    /// Requires test_data/futures-baseline-mbp10/ES.FUT/ with .dbn.zst files.
    /// Run with: cargo test -p ml-features --lib -- test_mbp10_to_imbalance_bars_real --ignored --nocapture
    #[test]
    #[ignore]
    fn test_mbp10_to_imbalance_bars_real() {
        // Resolve workspace root from CARGO_MANIFEST_DIR
        let manifest = std::env::var("CARGO_MANIFEST_DIR").unwrap();
        let workspace = std::path::Path::new(&manifest)
            .parent()  // crates/
            .and_then(|p| p.parent())  // repo root
            .expect("Failed to find workspace root");

        // Load threshold/alpha from TOML config (no hardcoded values)
        let toml_path = workspace.join("config/training/dqn-smoketest.toml");
        let toml_str = std::fs::read_to_string(&toml_path)
            .unwrap_or_else(|_| panic!("Failed to read {}", toml_path.display()));
        let toml: toml::Value = toml_str.parse()
            .unwrap_or_else(|e| panic!("Failed to parse TOML: {e}"));
        let training = toml.get("training").expect("[training] section missing");
        let threshold = training.get("imbalance_bar_threshold")
            .and_then(|v| v.as_float()).unwrap_or(1.0);
        let alpha = training.get("imbalance_bar_ewma_alpha")
            .and_then(|v| v.as_float()).unwrap_or(0.1);
        println!("Config: threshold={threshold}, alpha={alpha} (from dqn-smoketest.toml)");

        let mbp10_dir = workspace.join("test_data/futures-baseline-mbp10");
        assert!(mbp10_dir.exists(), "MBP-10 test data not found at {}", mbp10_dir.display());

        let bars = mbp10_to_imbalance_bars(&mbp10_dir, "ES.FUT", threshold, alpha)
            .expect("mbp10_to_imbalance_bars must succeed with real data");

        assert!(!bars.is_empty(), "Must produce imbalance bars from MBP-10 data");
        println!("Produced {} imbalance bars from MBP-10 data", bars.len());
        println!("  First bar: {:?}", bars.first().unwrap().timestamp);
        println!("  Last bar:  {:?}", bars.last().unwrap().timestamp);

        // Verify time-ordering
        for w in bars.windows(2) {
            assert!(
                w[0].timestamp <= w[1].timestamp,
                "Bars must be time-ordered: {:?} > {:?}",
                w[0].timestamp, w[1].timestamp
            );
        }

        // Verify reasonable bar properties
        for bar in &bars {
            assert!(bar.high >= bar.low, "high must >= low");
            assert!(bar.volume > 0.0, "volume must be positive");
            assert!(bar.close > 0.0, "close must be positive");
        }
        println!("All {} bars valid: time-ordered, positive volume, valid OHLC", bars.len());
    }
}
