//! Parquet data loader for real market data replay
//!
//! Provides efficient loading of real market data from Parquet files
//! for use in tests, benchmarks, backtesting, and ML training.
//!
//! # Architecture
//!
//! The loader provides two main APIs:
//! - **Bulk Loading**: Load entire Parquet file into memory (`load_all`)
//! - **Streaming**: Load events in batches for memory efficiency (`load_stream`)
//!
//! # Examples
//!
//! ## Load all events from Parquet file
//! ```no_run
//! use data::replay::ParquetDataLoader;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("test_data/real/parquet/ES.FUT_ohlcv-1m_2024-01-02.parquet");
//! let events = loader.load_all().await?;
//! println!("Loaded {} events", events.len());
//! # Ok(())
//! # }
//! ```
//!
//! ## Stream events for memory efficiency
//! ```no_run
//! use data::replay::ParquetDataLoader;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let loader = ParquetDataLoader::new("test_data/real/parquet/ES.FUT_ohlcv-1m_2024-01-02.parquet");
//! let mut stream = loader.load_stream().await?;
//!
//! while let Some(batch) = stream.next().await {
//!     let batch = batch?;
//!     println!("Processing batch of {} events", batch.len());
//! }
//! # Ok(())
//! # }
//! ```

use anyhow::{Context, Result};
use arrow::array::{as_primitive_array, Array, Float64Array, StringArray};
use arrow::datatypes::{TimestampNanosecondType, UInt64Type};
use arrow::record_batch::RecordBatch;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::fs::File;
use std::path::Path;
use trading_engine::types::metrics::{MarketDataEventType, ParquetMarketDataEvent};

/// Parquet data loader for real market data
///
/// Provides efficient loading capabilities for Parquet files containing
/// historical market data. Supports both bulk loading (all events at once)
/// and streaming (batch-by-batch processing).
pub struct ParquetDataLoader {
    file_path: String,
}

impl ParquetDataLoader {
    /// Create new loader for a Parquet file
    ///
    /// # Arguments
    /// * `file_path` - Path to the Parquet file to load
    ///
    /// # Examples
    /// ```
    /// use data::replay::ParquetDataLoader;
    ///
    /// let loader = ParquetDataLoader::new("market_data.parquet");
    /// ```
    pub fn new(file_path: impl AsRef<Path>) -> Self {
        Self {
            file_path: file_path.as_ref().to_string_lossy().to_string(),
        }
    }

    /// Load all events from Parquet file into memory
    ///
    /// Loads the entire Parquet file into a vector. Use this for smaller files
    /// or when you need random access to all events.
    ///
    /// # Returns
    /// * `Ok(Vec<ParquetMarketDataEvent>)` - All events from the file
    /// * `Err(_)` - If file cannot be read or parsed
    ///
    /// # Examples
    /// ```no_run
    /// # use data::replay::ParquetDataLoader;
    /// # async fn example() -> anyhow::Result<()> {
    /// let loader = ParquetDataLoader::new("market_data.parquet");
    /// let events = loader.load_all().await?;
    /// println!("Loaded {} events", events.len());
    /// # Ok(())
    /// # }
    /// ```
    pub async fn load_all(&self) -> Result<Vec<ParquetMarketDataEvent>> {
        // Read from file synchronously (Parquet doesn't have async reader)
        let file = File::open(&self.file_path)
            .with_context(|| format!("Failed to open Parquet file: {}", self.file_path))?;

        let builder = ParquetRecordBatchReaderBuilder::try_new(file)
            .context("Failed to create Parquet reader")?;

        let reader = builder.build().context("Failed to build Parquet reader")?;

        let mut events = Vec::new();

        for batch_result in reader {
            let batch = batch_result.context("Failed to read record batch")?;
            events.extend(Self::batch_to_events(&batch)?);
        }

        Ok(events)
    }

    /// Load events in streaming fashion (iterator-based)
    ///
    /// Returns a stream of event batches for memory-efficient processing.
    /// Each batch contains multiple events from a single RecordBatch.
    ///
    /// # Returns
    /// * `Ok(ParquetEventStream)` - Stream of event batches
    /// * `Err(_)` - If file cannot be opened
    ///
    /// # Examples
    /// ```no_run
    /// # use data::replay::ParquetDataLoader;
    /// # async fn example() -> anyhow::Result<()> {
    /// let loader = ParquetDataLoader::new("market_data.parquet");
    /// let mut stream = loader.load_stream().await?;
    ///
    /// while let Some(batch) = stream.next().await {
    ///     let events = batch?;
    ///     // Process events...
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn load_stream(&self) -> Result<ParquetEventStream> {
        let file = File::open(&self.file_path)
            .with_context(|| format!("Failed to open Parquet file: {}", self.file_path))?;

        let builder = ParquetRecordBatchReaderBuilder::try_new(file)
            .context("Failed to create Parquet reader")?;

        let reader = builder.build().context("Failed to build Parquet reader")?;

        Ok(ParquetEventStream { reader })
    }

    /// Convert Arrow RecordBatch to events
    ///
    /// Internal helper to convert Parquet's Arrow RecordBatch representation
    /// into our typed ParquetMarketDataEvent structures.
    ///
    /// # Arguments
    /// * `batch` - Arrow RecordBatch from Parquet reader
    ///
    /// # Returns
    /// * `Ok(Vec<ParquetMarketDataEvent>)` - Converted events
    /// * `Err(_)` - If batch schema is invalid or data cannot be parsed
    fn batch_to_events(batch: &RecordBatch) -> Result<Vec<ParquetMarketDataEvent>> {
        let num_rows = batch.num_rows();
        let mut events = Vec::with_capacity(num_rows);

        // Extract columns - all required fields
        let timestamp_ns = as_primitive_array::<TimestampNanosecondType>(
            batch
                .column_by_name("timestamp_ns")
                .context("Missing timestamp_ns column")?,
        );
        let symbol = batch
            .column_by_name("symbol")
            .context("Missing symbol column")?
            .as_any()
            .downcast_ref::<StringArray>()
            .context("Invalid symbol column type")?;
        let venue = batch
            .column_by_name("venue")
            .context("Missing venue column")?
            .as_any()
            .downcast_ref::<StringArray>()
            .context("Invalid venue column type")?;
        let event_type = batch
            .column_by_name("event_type")
            .context("Missing event_type column")?
            .as_any()
            .downcast_ref::<StringArray>()
            .context("Invalid event_type column type")?;
        let sequence = as_primitive_array::<UInt64Type>(
            batch
                .column_by_name("sequence")
                .context("Missing sequence column")?,
        );

        // Optional columns
        let price = batch
            .column_by_name("price")
            .and_then(|col| col.as_any().downcast_ref::<Float64Array>());
        let quantity = batch
            .column_by_name("quantity")
            .and_then(|col| col.as_any().downcast_ref::<Float64Array>());
        let latency_ns = batch
            .column_by_name("latency_ns")
            .and_then(|col| as_primitive_array::<UInt64Type>(col).into());
        let open = batch
            .column_by_name("open")
            .and_then(|col| col.as_any().downcast_ref::<Float64Array>());
        let high = batch
            .column_by_name("high")
            .and_then(|col| col.as_any().downcast_ref::<Float64Array>());
        let low = batch
            .column_by_name("low")
            .and_then(|col| col.as_any().downcast_ref::<Float64Array>());

        for i in 0..num_rows {
            let event_type_str = event_type.value(i);
            let parsed_event_type = match event_type_str {
                "Trade" => MarketDataEventType::Trade,
                "Quote" => MarketDataEventType::Quote,
                "OrderBookUpdate" => MarketDataEventType::OrderBookUpdate,
                "StatusUpdate" => MarketDataEventType::StatusUpdate,
                "Ohlcv" => MarketDataEventType::Ohlcv,
                _ => {
                    return Err(anyhow::anyhow!("Unknown event type: {}", event_type_str));
                },
            };

            events.push(ParquetMarketDataEvent {
                timestamp_ns: timestamp_ns.value(i) as u64,
                symbol: symbol.value(i).to_string(),
                venue: venue.value(i).to_string(),
                event_type: parsed_event_type,
                price: price.and_then(|arr| {
                    if arr.is_null(i) {
                        None
                    } else {
                        Some(arr.value(i))
                    }
                }),
                quantity: quantity.and_then(|arr| {
                    if arr.is_null(i) {
                        None
                    } else {
                        Some(arr.value(i))
                    }
                }),
                sequence: sequence.value(i),
                latency_ns: latency_ns.and_then(|arr| {
                    if arr.is_null(i) {
                        None
                    } else {
                        Some(arr.value(i))
                    }
                }),
                open: open.and_then(|arr| {
                    if arr.is_null(i) {
                        None
                    } else {
                        Some(arr.value(i))
                    }
                }),
                high: high.and_then(|arr| {
                    if arr.is_null(i) {
                        None
                    } else {
                        Some(arr.value(i))
                    }
                }),
                low: low.and_then(|arr| {
                    if arr.is_null(i) {
                        None
                    } else {
                        Some(arr.value(i))
                    }
                }),
            });
        }

        Ok(events)
    }
}

/// Streaming iterator over Parquet events
///
/// Provides batch-by-batch streaming of events for memory-efficient processing.
/// Each call to `next()` returns a batch of events from the Parquet file.
pub struct ParquetEventStream {
    reader: parquet::arrow::arrow_reader::ParquetRecordBatchReader,
}

impl ParquetEventStream {
    /// Get next batch of events
    ///
    /// Returns `None` when all batches have been read.
    ///
    /// # Returns
    /// * `Some(Ok(Vec<ParquetMarketDataEvent>))` - Next batch of events
    /// * `Some(Err(_))` - Error reading or parsing batch
    /// * `None` - End of stream
    pub async fn next(&mut self) -> Option<Result<Vec<ParquetMarketDataEvent>>> {
        // Note: Parquet reader is synchronous, but we provide async API for consistency
        match self.reader.next() {
            Some(Ok(batch)) => Some(ParquetDataLoader::batch_to_events(&batch)),
            Some(Err(e)) => Some(Err(anyhow::anyhow!("Failed to read batch: {}", e))),
            None => None,
        }
    }
}

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

    #[tokio::test]
    async fn test_parquet_loader_creation() {
        let loader = ParquetDataLoader::new("test.parquet");
        assert!(loader.file_path.contains("test.parquet"));
    }

    #[tokio::test]
    async fn test_parquet_loader_nonexistent_file() {
        let loader = ParquetDataLoader::new("nonexistent.parquet");
        let result = loader.load_all().await;
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Failed to open Parquet file"));
    }
}
