foxhunt/data/src/parquet_persistence.rs

//! # Parquet Market Data Persistence for Replay
//!
//! High-performance Parquet-based market data persistence system for backtesting
//! and trade replay capabilities in the Foxhunt HFT system.

use anyhow::{Context, Result};
use arrow::array::{Float64Array, StringArray, TimestampNanosecondArray, UInt64Array};
use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use arrow::record_batch::RecordBatch;
use parquet::arrow::ArrowWriter;
use parquet::file::properties::{EnabledStatistics, WriterProperties};
use serde::{Deserialize, Serialize};
use std::fs::File;
use std::path::Path;
use std::sync::Arc;
use tokio::sync::{mpsc, RwLock};
use tokio::time::{Duration, Instant};
use tracing::{debug, error, info, warn};

// Import the renamed Parquet-specific market data event
use common::metrics::ParquetMarketDataEvent as MarketDataEvent;

/// Parquet writer configuration
#[derive(Debug, Clone)]
pub struct ParquetConfig {
    pub base_path: String,
    pub batch_size: usize,
    pub flush_interval_ms: u64,
    pub compression: parquet::basic::Compression,
    pub enable_dictionary: bool,
    pub enable_statistics: EnabledStatistics,
}

impl Default for ParquetConfig {
    fn default() -> Self {
        Self {
            base_path: "./market_data".to_string(),
            batch_size: 10000,
            flush_interval_ms: 5000,
            compression: parquet::basic::Compression::SNAPPY,
            enable_dictionary: true,
            enable_statistics: EnabledStatistics::Page,
        }
    }
}

/// High-performance Parquet writer for market data
pub struct ParquetMarketDataWriter {
    config: ParquetConfig,
    buffer: Arc<RwLock<Vec<MarketDataEvent>>>,
    sender: mpsc::UnboundedSender<MarketDataEvent>,
    _writer_handle: tokio::task::JoinHandle<()>,
}

impl ParquetMarketDataWriter {
    /// Create new Parquet writer with background processing
    pub async fn new(config: ParquetConfig) -> Result<Self> {
        // Ensure base directory exists
        std::fs::create_dir_all(&config.base_path)
            .context("Failed to create Parquet base directory")?;

        let buffer = Arc::new(RwLock::new(Vec::with_capacity(config.batch_size * 2)));
        let (sender, receiver) = mpsc::unbounded_channel();

        let writer_handle =
            Self::spawn_writer_task(config.clone(), buffer.clone(), receiver).await?;

        Ok(Self {
            config,
            buffer,
            sender,
            _writer_handle: writer_handle,
        })
    }

    /// Record market data event (non-blocking)
    pub fn record(&self, event: MarketDataEvent) -> Result<()> {
        self.sender
            .send(event)
            .context("Failed to send market data event to writer")?;
        Ok(())
    }

    /// Get current buffer statistics
    pub async fn get_buffer_stats(&self) -> BufferStats {
        let buffer = self.buffer.read().await;
        BufferStats {
            buffered_events: buffer.len(),
            buffer_capacity: buffer.capacity(),
            utilization_percent: (buffer.len() as f64 / buffer.capacity() as f64) * 100.0,
        }
    }

    /// Spawn background writer task
    async fn spawn_writer_task(
        config: ParquetConfig,
        buffer: Arc<RwLock<Vec<MarketDataEvent>>>,
        mut receiver: mpsc::UnboundedReceiver<MarketDataEvent>,
    ) -> Result<tokio::task::JoinHandle<()>> {
        let handle = tokio::spawn(async move {
            let mut flush_interval =
                tokio::time::interval(Duration::from_millis(config.flush_interval_ms));
            let mut last_flush = Instant::now();

            loop {
                tokio::select! {
                    // Handle incoming events
                    event = receiver.recv() => {
                        match event {
                            Some(event) => {
                                let mut buffer_guard = buffer.write().await;
                                buffer_guard.push(event);

                                // Check if we should flush based on batch size
                                if buffer_guard.len() >= config.batch_size {
                                    let events = buffer_guard.drain(..).collect();
                                    drop(buffer_guard);

                                    if let Err(e) = Self::write_batch_to_parquet(&config, events).await {
                                        error!("Failed to write Parquet batch: {}", e);
                                    }
                                    last_flush = Instant::now();
                                }
                            }
                            None => {
                                info!("Parquet writer channel closed, shutting down");
                                break;
                            }
                        }
                    }

                    // Handle periodic flush
                    _ = flush_interval.tick() => {
                        if last_flush.elapsed() >= Duration::from_millis(config.flush_interval_ms) {
                            let mut buffer_guard = buffer.write().await;
                            if !buffer_guard.is_empty() {
                                let events = buffer_guard.drain(..).collect();
                                drop(buffer_guard);

                                if let Err(e) = Self::write_batch_to_parquet(&config, events).await {
                                    error!("Failed to write Parquet batch on flush: {}", e);
                                }
                                last_flush = Instant::now();
                            }
                        }
                    }
                }
            }

            // Final flush on shutdown
            let mut buffer_guard = buffer.write().await;
            if !buffer_guard.is_empty() {
                let events = buffer_guard.drain(..).collect();
                drop(buffer_guard);

                if let Err(e) = Self::write_batch_to_parquet(&config, events).await {
                    error!("Failed to write final Parquet batch: {}", e);
                }
            }
        });

        Ok(handle)
    }

    /// Write batch of events to Parquet file
    async fn write_batch_to_parquet(
        config: &ParquetConfig,
        events: Vec<MarketDataEvent>,
    ) -> Result<()> {
        if events.is_empty() {
            return Ok(());
        }

        let start_time = Instant::now();

        // Generate filename with timestamp
        let timestamp = events[0].timestamp_ns;
        let date = chrono::DateTime::from_timestamp_nanos(timestamp as i64).format("%Y%m%d_%H%M%S");
        let filename = format!(
            "market_data_{}_{}.parquet",
            date,
            uuid::Uuid::new_v4().simple()
        );
        let filepath = Path::new(&config.base_path).join(filename);

        // Create Arrow schema
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "timestamp_ns",
                DataType::Timestamp(TimeUnit::Nanosecond, None),
                false,
            ),
            Field::new("symbol", DataType::Utf8, false),
            Field::new("venue", DataType::Utf8, false),
            Field::new("event_type", DataType::Utf8, false),
            Field::new("price", DataType::Float64, true),
            Field::new("quantity", DataType::Float64, true),
            Field::new("bid_price", DataType::Float64, true),
            Field::new("ask_price", DataType::Float64, true),
            Field::new("bid_size", DataType::Float64, true),
            Field::new("ask_size", DataType::Float64, true),
            Field::new("sequence", DataType::UInt64, false),
            Field::new("latency_ns", DataType::UInt64, true),
        ]));

        // Convert events to Arrow arrays
        let record_batch = Self::events_to_record_batch(&schema, events)?;

        // Write to Parquet file
        let file = File::create(&filepath)
            .with_context(|| format!("Failed to create Parquet file: {:?}", filepath))?;

        let props = WriterProperties::builder()
            .set_compression(config.compression)
            .set_dictionary_enabled(config.enable_dictionary)
            .set_statistics_enabled(config.enable_statistics)
            .build();

        let mut writer = ArrowWriter::try_new(file, schema, Some(props))
            .context("Failed to create Arrow writer")?;

        writer
            .write(&record_batch)
            .context("Failed to write record batch")?;

        writer.close().context("Failed to close Arrow writer")?;

        let duration = start_time.elapsed();
        let events_count = record_batch.num_rows();

        debug!(
            "Wrote {} events to Parquet file {:?} in {:?}",
            events_count, filepath, duration
        );

        // Update metrics
        let duration_us: u64 = duration.as_micros().try_into().unwrap_or(0);
        if duration_us > 0 {
            common::metrics::LATENCY_HISTOGRAMS
                .with_label_values(&["parquet_write", "data_service"])
                .observe(duration_us as f64 / 1_000_000.0);
        }

        common::metrics::THROUGHPUT_COUNTERS
            .with_label_values(&["parquet_events", "data_service"])
            .inc_by(events_count as u64);

        Ok(())
    }

    /// Convert events to Arrow RecordBatch
    fn events_to_record_batch(
        schema: &Arc<Schema>,
        events: Vec<MarketDataEvent>,
    ) -> Result<RecordBatch> {
        let len = events.len();

        // Extract data into separate vectors
        let mut timestamps = Vec::with_capacity(len);
        let mut symbols = Vec::with_capacity(len);
        let mut venues = Vec::with_capacity(len);
        let mut event_types = Vec::with_capacity(len);
        let mut prices = Vec::with_capacity(len);
        let mut quantities = Vec::with_capacity(len);
        let mut bid_prices = Vec::with_capacity(len);
        let mut ask_prices = Vec::with_capacity(len);
        let mut bid_sizes = Vec::with_capacity(len);
        let mut ask_sizes = Vec::with_capacity(len);
        let mut sequences = Vec::with_capacity(len);
        let mut latencies = Vec::with_capacity(len);

        for event in events {
            timestamps.push(Some(event.timestamp_ns as i64));
            symbols.push(Some(event.symbol));
            venues.push(Some(event.venue));
            event_types.push(Some(event.event_type));
            prices.push(event.price);
            quantities.push(event.quantity);
            bid_prices.push(event.bid_price);
            ask_prices.push(event.ask_price);
            bid_sizes.push(event.bid_size);
            ask_sizes.push(event.ask_size);
            sequences.push(event.sequence);
            latencies.push(event.latency_ns);
        }

        // Create Arrow arrays
        let timestamp_array = TimestampNanosecondArray::from(timestamps);
        let symbol_array = StringArray::from(symbols);
        let venue_array = StringArray::from(venues);
        let event_type_array = StringArray::from(event_types);
        let price_array = Float64Array::from(prices);
        let quantity_array = Float64Array::from(quantities);
        let bid_price_array = Float64Array::from(bid_prices);
        let ask_price_array = Float64Array::from(ask_prices);
        let bid_size_array = Float64Array::from(bid_sizes);
        let ask_size_array = Float64Array::from(ask_sizes);
        let sequence_array = UInt64Array::from(sequences);
        let latency_array = UInt64Array::from(latencies);

        // Create record batch
        RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(timestamp_array),
                Arc::new(symbol_array),
                Arc::new(venue_array),
                Arc::new(event_type_array),
                Arc::new(price_array),
                Arc::new(quantity_array),
                Arc::new(bid_price_array),
                Arc::new(ask_price_array),
                Arc::new(bid_size_array),
                Arc::new(ask_size_array),
                Arc::new(sequence_array),
                Arc::new(latency_array),
            ],
        )
        .context("Failed to create Arrow RecordBatch")
    }
}

/// Buffer statistics for monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BufferStats {
    pub buffered_events: usize,
    pub buffer_capacity: usize,
    pub utilization_percent: f64,
}

/// Parquet reader for market data replay
pub struct ParquetMarketDataReader {
    base_path: String,
}

impl ParquetMarketDataReader {
    pub fn new(base_path: String) -> Self {
        Self { base_path }
    }

    /// List available Parquet files for replay
    pub async fn list_available_files(&self) -> Result<Vec<String>> {
        let mut files = Vec::new();
        let entries =
            std::fs::read_dir(&self.base_path).context("Failed to read Parquet directory")?;

        for entry in entries {
            let entry = entry.context("Failed to read directory entry")?;
            let path = entry.path();

            if path.is_file() && path.extension().map_or(false, |ext| ext == "parquet") {
                if let Some(filename) = path.file_name().and_then(|f| f.to_str()) {
                    files.push(filename.to_string());
                }
            }
        }

        files.sort();
        Ok(files)
    }

    /// Read market data from Parquet file for replay
    pub async fn read_file(&self, filename: &str) -> Result<Vec<MarketDataEvent>> {
        let filepath = Path::new(&self.base_path).join(filename);

        // This would be implemented using parquet::arrow::async_reader
        // For now, return placeholder
        warn!("Parquet reader not fully implemented yet: {:?}", filepath);
        Ok(Vec::new())
    }
}

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

    #[tokio::test]
    async fn test_parquet_writer_creation() {
        let temp_dir = tempdir().unwrap();
        let config = ParquetConfig {
            base_path: temp_dir.path().to_string_lossy().to_string(),
            ..Default::default()
        };

        let writer = ParquetMarketDataWriter::new(config).await;
        assert!(writer.is_ok());
    }

    #[tokio::test]
    async fn test_market_data_event_recording() {
        let temp_dir = tempdir().unwrap();
        let config = ParquetConfig {
            base_path: temp_dir.path().to_string_lossy().to_string(),
            batch_size: 2, // Small batch for testing
            ..Default::default()
        };

        let writer = ParquetMarketDataWriter::new(config).await.unwrap();

        let event = MarketDataEvent {
            timestamp_ns: 1234567890000000000,
            symbol: "BTCUSD".to_string(),
            venue: "binance".to_string(),
            event_type: "trade".to_string(),
            price: Some(50000.0),
            quantity: Some(0.1),
            bid_price: None,
            ask_price: None,
            bid_size: None,
            ask_size: None,
            sequence: 1,
            latency_ns: Some(1000),
        };

        let result = writer.record(event);
        assert!(result.is_ok());

        // Give some time for background processing
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
}