//! TLOB Data Loader for Level 2 Order Book Training
//!
//! Loads DataBento MBP-10 (Market By Price, 10 levels) data for TLOB Transformer training.
//! Provides tick-by-tick order book snapshots with 10 bid and 10 ask price levels.
//!
//! ## Features
//!
//! - Loads MBP-10 DBN files (Level 2 order book data)
//! - Extracts 10 bid/ask price levels per snapshot
//! - Creates fixed-length sequences for transformer training
//! - Integrates with TLOBFeatureExtractor (51 features)
//! - Handles temporal ordering and sequence creation
//!
//! ## Usage
//!
//! ```no_run
//! use ml::data_loaders::TLOBDataLoader;
//! use NativeDevice;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = TLOBDataLoader::new(128, 51).await?;
//! let (train_data, val_data) = loader
//!     .load_sequences("test_data/real/databento/l2_order_book", 0.9)
//!     .await?;
//!
//! println!("Loaded {} training sequences", train_data.len());
//! # Ok(())
//! # }
//! ```

use anyhow::{Context, Result};
use ml_core::native_types::NativeDevice;
use ml_core::cuda_autograd::GpuTensor;
use dbn::decode::{DbnDecoder, DbnMetadata, DecodeRecordRef};
use dbn::RecordRefEnum;
use std::collections::HashMap;
use std::fs::File;
use std::io::BufReader;
use std::path::Path;
use tokio::fs;
use tracing::{debug, info, warn};

use crate::tlob::{TLOBFeatureExtractor, TLOBInputFeatures, TLOB_FEATURE_COUNT};

/// TLOB data loader for Level 2 order book training
pub struct TLOBDataLoader {
    /// Target sequence length (number of order book snapshots)
    seq_len: usize,

    /// Feature dimension (should be 51 for TLOB)
    feature_dim: usize,

    /// NativeDevice for tensor creation
    device: NativeDevice,

    /// TLOB feature extractor
    feature_extractor: TLOBFeatureExtractor,
}

impl std::fmt::Debug for TLOBDataLoader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TLOBDataLoader")
            .field("seq_len", &self.seq_len)
            .field("feature_dim", &self.feature_dim)
            .finish_non_exhaustive()
    }
}

/// Order book snapshot from MBP-10 data
#[derive(Debug, Clone)]
pub struct OrderBookSnapshot {
    pub timestamp: u64,
    pub symbol: String,
    pub bid_levels: Vec<i64>,  // 10 bid prices (1e-9 fixed-point)
    pub ask_levels: Vec<i64>,  // 10 ask prices (1e-9 fixed-point)
    pub bid_volumes: Vec<i64>, // 10 bid sizes
    pub ask_volumes: Vec<i64>, // 10 ask sizes
    pub last_price: i64,       // Most recent trade price
    pub volume: i64,           // Cumulative volume
}

impl TLOBDataLoader {
    /// Create new TLOB data loader
    ///
    /// # Arguments
    /// * `seq_len` - Target sequence length (64-256 recommended for transformers)
    /// * `feature_dim` - Feature dimension (should be 51 for TLOB)
    ///
    /// # Returns
    /// Configured loader ready to process MBP-10 DBN files
    pub async fn new(seq_len: usize, feature_dim: usize) -> Result<Self> {
        if feature_dim != TLOB_FEATURE_COUNT {
            warn!(
                "Feature dimension {} does not match TLOB_FEATURE_COUNT ({}). Using {}.",
                feature_dim, TLOB_FEATURE_COUNT, feature_dim
            );
        }

        let device = NativeDevice::Cuda(0);

        let feature_extractor = TLOBFeatureExtractor::new()
            .map_err(|e| anyhow::anyhow!("Failed to create TLOB feature extractor: {}", e))?;

        info!(
            "TLOB data loader initialized (seq_len={}, feature_dim={}, device={:?})",
            seq_len, feature_dim, device
        );

        Ok(Self {
            seq_len,
            feature_dim,
            device,
            feature_extractor,
        })
    }

    /// Load sequences from a directory of MBP-10 DBN files
    ///
    /// # Arguments
    /// * `dbn_dir` - Directory containing .dbn files with MBP-10 data
    /// * `train_split` - Fraction of data for training (0.0-1.0)
    ///
    /// # Returns
    /// Tuple of (train_sequences, val_sequences) as (input, target) pairs
    pub async fn load_sequences<P: AsRef<Path>>(
        &mut self,
        dbn_dir: P,
        train_split: f64,
    ) -> Result<(Vec<(GpuTensor, GpuTensor)>, Vec<(GpuTensor, GpuTensor)>)> {
        let path = dbn_dir.as_ref();
        info!("Loading MBP-10 sequences from: {:?}", path);

        // Find all .dbn files
        let mut dbn_files = Vec::new();
        let mut entries = fs::read_dir(path)
            .await
            .with_context(|| format!("Failed to read directory: {:?}", path))?;

        while let Some(entry) = entries.next_entry().await? {
            let path = entry.path();
            if path.extension().and_then(|s| s.to_str()) == Some("dbn") {
                dbn_files.push(path);
            }
        }

        dbn_files.sort();
        info!("Found {} MBP-10 DBN files", dbn_files.len());

        if dbn_files.is_empty() {
            return Err(anyhow::anyhow!("No DBN files found in {:?}", path));
        }

        // Load all snapshots from all files and group by symbol
        let mut symbol_snapshots: HashMap<String, Vec<OrderBookSnapshot>> = HashMap::new();

        for file_path in &dbn_files {
            info!("Processing: {:?}", file_path);
            let snapshots = self.load_file(file_path).await?;

            // Group by symbol for temporal ordering
            for snapshot in snapshots {
                symbol_snapshots
                    .entry(snapshot.symbol.clone())
                    .or_default()
                    .push(snapshot);
            }
        }

        info!(
            "Loaded order book snapshots for {} symbols",
            symbol_snapshots.len()
        );

        // Create sequences from each symbol's snapshots
        let mut all_sequences = Vec::new();

        for (symbol, snapshots) in symbol_snapshots {
            if snapshots.len() < self.seq_len + 1 {
                warn!(
                    "Skipping {}: only {} snapshots (need {})",
                    symbol,
                    snapshots.len(),
                    self.seq_len + 1
                );
                continue;
            }

            let sequences = self.create_sequences(&snapshots)?;
            all_sequences.extend(sequences);
            debug!("Created {} sequences from {}", all_sequences.len(), symbol);
        }

        info!("Created {} total sequences", all_sequences.len());

        // Split into train/val
        let split_idx = (all_sequences.len() as f64 * train_split) as usize;
        let val_data = all_sequences.split_off(split_idx);
        let train_data = all_sequences;

        info!(
            "Split: {} training, {} validation",
            train_data.len(),
            val_data.len()
        );

        Ok((train_data, val_data))
    }

    /// Load order book snapshots from a single MBP-10 DBN file
    async fn load_file<P: AsRef<Path>>(&self, path: P) -> Result<Vec<OrderBookSnapshot>> {
        let path = path.as_ref();

        // Open file and create DBN decoder
        let file = File::open(path).with_context(|| format!("Failed to open: {:?}", path))?;
        let reader = BufReader::new(file);

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

        // Read metadata
        let metadata = decoder.metadata();
        let symbol = metadata
            .symbols
            .first()
            .cloned()
            .unwrap_or_else(|| "UNKNOWN".to_owned());

        debug!(
            "MBP-10 file metadata: dataset={:?}, schema={:?}, symbol={}",
            metadata.dataset, metadata.schema, symbol
        );

        // Decode all MBP-10 records
        let mut snapshots = Vec::new();
        let mut mbp10_count = 0;
        let mut last_price = 0_i64;
        let mut cumulative_volume = 0_i64;

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

                    match record_enum {
                        RecordRefEnum::Mbp10(mbp) => {
                            mbp10_count += 1;

                            // Extract 10 bid/ask levels
                            let mut bid_levels = Vec::with_capacity(10);
                            let mut ask_levels = Vec::with_capacity(10);
                            let mut bid_volumes = Vec::with_capacity(10);
                            let mut ask_volumes = Vec::with_capacity(10);

                            for level in &mbp.levels {
                                bid_levels.push(level.bid_px);
                                ask_levels.push(level.ask_px);
                                bid_volumes.push(level.bid_sz as i64);
                                ask_volumes.push(level.ask_sz as i64);
                            }

                            // Update last price (use mid price from level 0)
                            if !bid_levels.is_empty() && !ask_levels.is_empty() {
                                last_price = (bid_levels[0] + ask_levels[0]) / 2;
                            }

                            // Update cumulative volume
                            let snapshot_volume: i64 =
                                bid_volumes.iter().sum::<i64>() + ask_volumes.iter().sum::<i64>();
                            cumulative_volume += snapshot_volume;

                            snapshots.push(OrderBookSnapshot {
                                timestamp: mbp.hd.ts_event,
                                symbol: symbol.clone(),
                                bid_levels,
                                ask_levels,
                                bid_volumes,
                                ask_volumes,
                                last_price,
                                volume: cumulative_volume,
                            });

                            // Log first few snapshots for validation
                            if mbp10_count <= 3 {
                                let bid_f64 = mbp.levels[0].bid_px as f64 * 1e-9;
                                let ask_f64 = mbp.levels[0].ask_px as f64 * 1e-9;
                                debug!(
                                    "Snapshot #{}: timestamp={}, bid={:.2}, ask={:.2}, spread={:.4}",
                                    mbp10_count,
                                    mbp.hd.ts_event,
                                    bid_f64,
                                    ask_f64,
                                    ask_f64 - bid_f64
                                );
                            }
                        },
                        RecordRefEnum::Mbo(_)
                        | RecordRefEnum::Trade(_)
                        | RecordRefEnum::Mbp1(_)
                        | RecordRefEnum::Ohlcv(_)
                        | RecordRefEnum::Status(_)
                        | RecordRefEnum::InstrumentDef(_)
                        | RecordRefEnum::Imbalance(_)
                        | RecordRefEnum::Stat(_)
                        | RecordRefEnum::Error(_)
                        | RecordRefEnum::SymbolMapping(_)
                        | RecordRefEnum::System(_)
                        | RecordRefEnum::Cmbp1(_)
                        | RecordRefEnum::Bbo(_)
                        | RecordRefEnum::Cbbo(_) => {
                            // Skip other record types
                        },
                    }
                },
                Ok(None) => {
                    // End of stream
                    break;
                },
                Err(e) => {
                    return Err(anyhow::anyhow!("Failed to decode record: {}", e));
                },
            }
        }

        info!(
            "Loaded {} order book snapshots from {:?}",
            mbp10_count,
            path.file_name().unwrap_or_default()
        );

        Ok(snapshots)
    }

    /// Create sequences from order book snapshot list
    fn create_sequences(&self, snapshots: &[OrderBookSnapshot]) -> Result<Vec<(GpuTensor, GpuTensor)>> {
        let mut sequences = Vec::new();

        // Sliding window over snapshots
        for i in 0..snapshots.len().saturating_sub(self.seq_len) {
            let window = &snapshots[i..i + self.seq_len + 1];

            // Extract features for seq_len steps
            let mut features = Vec::with_capacity(self.seq_len * self.feature_dim);

            for snapshot in &window[..self.seq_len] {
                let snapshot_features = self.extract_features(snapshot)?;

                // Pad or truncate to feature_dim
                for j in 0..self.feature_dim {
                    if j < snapshot_features.len() {
                        features.push(snapshot_features[j]);
                    } else {
                        features.push(0.0); // Zero padding
                    }
                }
            }

            // Target is next timestep (autoregressive)
            let target_snapshot = &window[self.seq_len];
            let target_features = self.extract_features(target_snapshot)?;
            let mut target = vec![0.0; self.feature_dim];
            for j in 0..self.feature_dim.min(target_features.len()) {
                target[j] = target_features[j];
            }

            // Create tensors
            let ordinal = self.device.cuda_ordinal().unwrap_or(0);
            let ml_dev = ml_core::device::MlDevice::cuda(ordinal)
                .map_err(|e| anyhow::anyhow!("CUDA device for TLOB: {e}"))?;
            let stream = ml_dev.cuda_stream()
                .map_err(|e| anyhow::anyhow!("CUDA stream for TLOB: {e}"))?;

            let input =
                GpuTensor::from_host(&features, vec![self.seq_len, self.feature_dim], stream)?;

            let target_tensor = GpuTensor::from_host(&target, vec![1, self.feature_dim], stream)?;

            sequences.push((input, target_tensor));
        }

        Ok(sequences)
    }

    /// Extract 51 TLOB features from an order book snapshot
    fn extract_features(&self, snapshot: &OrderBookSnapshot) -> Result<Vec<f32>> {
        // Convert to TLOBFeatures format
        let tlob_features = TLOBInputFeatures::new(
            snapshot.timestamp,
            snapshot.symbol.clone(),
            snapshot.bid_levels.clone(),
            snapshot.ask_levels.clone(),
            snapshot.bid_volumes.clone(),
            snapshot.ask_volumes.clone(),
            snapshot.last_price,
            snapshot.volume,
            0.0,    // volatility (computed by feature extractor)
            0.0,    // momentum (computed by feature extractor)
            vec![], // microstructure features (computed by feature extractor)
        )
        .map_err(|e| anyhow::anyhow!("Failed to create TLOB features: {}", e))?;

        // Extract 51 features using TLOB feature extractor
        let feature_vector = self
            .feature_extractor
            .extract(&tlob_features)
            .map_err(|e| anyhow::anyhow!("Failed to extract features: {}", e))?;

        // Convert to f32 vector
        let features: Vec<f32> = feature_vector.values.iter().map(|&v| v as f32).collect();

        // Validate feature count
        if features.len() != TLOB_FEATURE_COUNT {
            warn!(
                "Expected {} features, got {}. Padding/truncating.",
                TLOB_FEATURE_COUNT,
                features.len()
            );
        }

        Ok(features)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_loader_creation() {
        let loader = TLOBDataLoader::new(128, 51).await;
        assert!(loader.is_ok());
    }

    #[tokio::test]
    async fn test_feature_dimension_validation() {
        // Should warn but still create loader
        let loader = TLOBDataLoader::new(128, 64).await;
        assert!(loader.is_ok());
    }

    #[test]
    fn test_order_book_snapshot_creation() {
        let snapshot = OrderBookSnapshot {
            timestamp: 1234567890,
            symbol: "ES.FUT".to_owned(),
            bid_levels: vec![4500; 10],
            ask_levels: vec![4501; 10],
            bid_volumes: vec![100; 10],
            ask_volumes: vec![100; 10],
            last_price: 4500,
            volume: 1000,
        };

        assert_eq!(snapshot.bid_levels.len(), 10);
        assert_eq!(snapshot.ask_levels.len(), 10);
    }
}
