use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock};

use arrow::array::{BooleanArray, TimestampNanosecondArray, UInt32Array};
use arrow::compute::filter_record_batch;
use arrow::record_batch::RecordBatch;
use bt_kernel::{concat_batches, BtError, Result, SymbolId, TimeRange};
use bytes::Bytes;
use memmap2::Mmap;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::file::metadata::ParquetMetaData;
use parquet::file::reader::{FileReader, SerializedFileReader};
use parquet::file::statistics::Statistics;
use rayon::prelude::*;

use crate::metadata::SqliteMetadataStore;
use crate::schema::canonical_bars_1m_schema;
use crate::store::DataStore;
use crate::version::DatasetVersion;
use crate::WriteOptions;

type CacheKey = (SymbolId, i64, i64);

/// A DataStore backed by a single memory-mapped mega parquet file.
///
/// The mega file contains all symbols sorted by (symbol_id, timestamp).
/// Row group statistics allow skipping irrelevant row groups, making
/// single-symbol queries fast even from a multi-GB file.
pub struct MegaDataStore {
    _path: PathBuf,
    data: Bytes,
    parquet_metadata: Arc<ParquetMetaData>,
    /// Pre-computed: column index of "symbol_id" in parquet schema.
    sym_col_idx: Option<usize>,
    /// Pre-computed: column index of "timestamp" in parquet schema.
    ts_col_idx: Option<usize>,
    metadata: SqliteMetadataStore,
    cache: RwLock<HashMap<CacheKey, RecordBatch>>,
}

impl MegaDataStore {
    pub fn new(
        mega_path: impl AsRef<Path>,
        metadata_path: impl AsRef<Path>,
    ) -> Result<Self> {
        let mega_path = mega_path.as_ref().to_path_buf();
        tracing::info!(path = %mega_path.display(), "opening mega data store (mmap)");

        let file = File::open(&mega_path)
            .map_err(|e| BtError::Data(format!("opening mega file {}: {e}", mega_path.display())))?;

        // Safety: the file must not be modified while mmap'd.
        // Bytes::from_owner takes ownership of the Mmap — zero-copy, no heap allocation.
        let mmap = unsafe { Mmap::map(&file) }
            .map_err(|e| BtError::Data(format!("mmap mega file: {e}")))?;

        let data = Bytes::from_owner(mmap);

        // Pre-read parquet metadata (footer) — cheap, gives us row group stats
        let reader = SerializedFileReader::new(data.clone())
            .map_err(|e| BtError::Data(format!("reading mega parquet metadata: {e}")))?;
        let parquet_metadata = reader.metadata().clone();

        let metadata = SqliteMetadataStore::open(metadata_path)?;

        // Pre-compute column indices for fast row group pruning
        let schema = parquet_metadata.file_metadata().schema_descr();
        let sym_col_idx = schema
            .columns()
            .iter()
            .position(|c| c.name() == "symbol_id");
        let ts_col_idx = schema
            .columns()
            .iter()
            .position(|c| c.name() == "timestamp");

        let n_rg = parquet_metadata.num_row_groups();
        let total_rows: i64 = (0..n_rg)
            .map(|i| parquet_metadata.row_group(i).num_rows())
            .sum();
        tracing::info!(
            row_groups = n_rg,
            total_rows,
            size_mb = data.len() / (1024 * 1024),
            "mega store ready"
        );

        Ok(Self {
            _path: mega_path,
            data,
            parquet_metadata: Arc::new(parquet_metadata),
            sym_col_idx,
            ts_col_idx,
            metadata,
            cache: RwLock::new(HashMap::new()),
        })
    }

    pub fn metadata_store(&self) -> &SqliteMetadataStore {
        &self.metadata
    }

    /// Find row group indices whose statistics overlap with (symbol_id, time_range).
    fn relevant_row_groups(&self, symbol: SymbolId, range: TimeRange) -> Vec<usize> {
        let meta = &self.parquet_metadata;
        let n_rg = meta.num_row_groups();
        let mut groups = Vec::new();

        for rg_idx in 0..n_rg {
            if self.rg_matches_symbol_and_time(rg_idx, symbol, range) {
                groups.push(rg_idx);
            }
        }

        groups
    }

    /// Check if a row group's statistics overlap with (symbol_id, time_range).
    #[inline]
    fn rg_matches_symbol_and_time(
        &self,
        rg_idx: usize,
        symbol: SymbolId,
        range: TimeRange,
    ) -> bool {
        let rg = self.parquet_metadata.row_group(rg_idx);

        if let Some(col_idx) = self.sym_col_idx {
            if let Some(Statistics::Int32(s)) = rg.column(col_idx).statistics() {
                let sid = symbol.0 as i32;
                if let (Some(&min), Some(&max)) = (s.min_opt(), s.max_opt()) {
                    if sid < min || sid > max {
                        return false;
                    }
                }
            }
        }

        if let Some(col_idx) = self.ts_col_idx {
            if let Some(Statistics::Int64(s)) = rg.column(col_idx).statistics() {
                if let (Some(&rg_min), Some(&rg_max)) = (s.min_opt(), s.max_opt()) {
                    if rg_max < range.start.0 || rg_min >= range.end.0 {
                        return false;
                    }
                }
            }
        }

        true
    }

    /// Read specific row groups from the mega file in a single reader pass,
    /// filtering to the given symbol and time range.
    fn read_row_groups(
        &self,
        row_groups: Vec<usize>,
        symbol: SymbolId,
        range: TimeRange,
    ) -> Result<RecordBatch> {
        if row_groups.is_empty() {
            return Ok(RecordBatch::new_empty(canonical_bars_1m_schema()));
        }

        // Single reader for ALL relevant row groups — parses footer only once.
        let reader = ParquetRecordBatchReaderBuilder::try_new(self.data.clone())
            .map_err(|e| BtError::Data(format!("mega parquet reader: {e}")))?
            .with_batch_size(262_144)
            .with_row_groups(row_groups)
            .build()
            .map_err(|e| BtError::Data(format!("mega parquet reader build: {e}")))?;

        let mut all_batches = Vec::new();

        for batch_result in reader {
            let batch = batch_result
                .map_err(|e| BtError::Data(format!("reading mega row group: {e}")))?;
            let filtered = filter_batch_by_symbol_and_time(&batch, symbol, range)?;
            if filtered.num_rows() > 0 {
                all_batches.push(filtered);
            }
        }

        if all_batches.is_empty() {
            return Ok(RecordBatch::new_empty(canonical_bars_1m_schema()));
        }

        concat_batches(canonical_bars_1m_schema(), &all_batches)
    }

    /// Preload all symbols in the given time range using parallel row group reading.
    /// Each row group is decompressed in parallel via rayon, then results are
    /// partitioned by symbol_id and cached.
    pub fn preload(
        &self,
        symbols: &[SymbolId],
        range: TimeRange,
    ) -> Result<()> {
        let meta = &self.parquet_metadata;
        let n_rg = meta.num_row_groups();

        // Build a set of requested symbol IDs
        let sym_set: std::collections::HashSet<u32> =
            symbols.iter().map(|s| s.0).collect();

        // Find ALL row groups that overlap with the time range for ANY requested symbol
        let mut row_groups = Vec::new();
        for rg_idx in 0..n_rg {
            let rg = meta.row_group(rg_idx);

            // Check symbol_id statistics
            if let Some(col_idx) = self.sym_col_idx {
                if let Some(Statistics::Int32(s)) = rg.column(col_idx).statistics() {
                    if let (Some(&min), Some(&max)) = (s.min_opt(), s.max_opt()) {
                        let any_match = sym_set.iter().any(|&sid| {
                            let sid_i32 = sid as i32;
                            sid_i32 >= min && sid_i32 <= max
                        });
                        if !any_match {
                            continue;
                        }
                    }
                }
            }

            // Check timestamp statistics
            if let Some(col_idx) = self.ts_col_idx {
                if let Some(Statistics::Int64(s)) = rg.column(col_idx).statistics() {
                    if let (Some(&rg_min), Some(&rg_max)) = (s.min_opt(), s.max_opt()) {
                        if rg_max < range.start.0 || rg_min >= range.end.0 {
                            continue;
                        }
                    }
                }
            }

            row_groups.push(rg_idx);
        }

        tracing::info!(
            total_rg = n_rg,
            selected_rg = row_groups.len(),
            symbols = symbols.len(),
            "mega store: parallel preload"
        );

        if row_groups.is_empty() {
            return Ok(());
        }

        let start = range.start.0;
        let end = range.end.0;

        // Parallel: each row group decompressed independently via rayon.
        // Returns Vec<(symbol_id, RecordBatch)> per row group.
        let per_rg_results: Vec<Vec<(u32, RecordBatch)>> = row_groups
            .par_iter()
            .map(|&rg_idx| {
                let reader = ParquetRecordBatchReaderBuilder::try_new(self.data.clone())
                    .map_err(|e| BtError::Data(format!("mega parquet reader: {e}")))?
                    .with_batch_size(500_000)
                    .with_row_groups(vec![rg_idx])
                    .build()
                    .map_err(|e| BtError::Data(format!("mega parquet reader build: {e}")))?;

                let mut rg_results: Vec<(u32, RecordBatch)> = Vec::new();

                for batch_result in reader {
                    let batch = batch_result
                        .map_err(|e| BtError::Data(format!("reading mega: {e}")))?;

                    let ts_col = batch
                        .column_by_name("timestamp")
                        .and_then(|c| c.as_any().downcast_ref::<TimestampNanosecondArray>())
                        .ok_or_else(|| BtError::Data("missing timestamp column".to_string()))?;
                    let sym_col = batch
                        .column_by_name("symbol_id")
                        .and_then(|c| c.as_any().downcast_ref::<UInt32Array>())
                        .ok_or_else(|| BtError::Data("missing symbol_id column".to_string()))?;

                    // Fast path: single-symbol batch (common since data is sorted by symbol_id)
                    let first_sym = sym_col.value(0);
                    let last_sym = sym_col.value(sym_col.len() - 1);

                    if first_sym == last_sym && sym_set.contains(&first_sym) {
                        let first_ts = ts_col.value(0);
                        let last_ts = ts_col.value(ts_col.len() - 1);

                        if first_ts >= start && last_ts < end {
                            // Entire batch in range — zero filter
                            rg_results.push((first_sym, batch));
                        } else {
                            // Time filter only
                            let mask = BooleanArray::from(
                                ts_col
                                    .values()
                                    .iter()
                                    .map(|&ts| ts >= start && ts < end)
                                    .collect::<Vec<bool>>(),
                            );
                            let filtered = filter_record_batch(&batch, &mask)?;
                            if filtered.num_rows() > 0 {
                                rg_results.push((first_sym, filtered));
                            }
                        }
                    } else {
                        // Mixed-symbol batch — split per symbol
                        for &sid in &sym_set {
                            let mask = BooleanArray::from(
                                ts_col
                                    .values()
                                    .iter()
                                    .zip(sym_col.values().iter())
                                    .map(|(&ts, &s)| s == sid && ts >= start && ts < end)
                                    .collect::<Vec<bool>>(),
                            );
                            let filtered = filter_record_batch(&batch, &mask)?;
                            if filtered.num_rows() > 0 {
                                rg_results.push((sid, filtered));
                            }
                        }
                    }
                }

                Ok(rg_results)
            })
            .collect::<Result<_>>()?;

        // Merge parallel results per symbol
        let mut per_symbol: HashMap<u32, Vec<RecordBatch>> = HashMap::new();
        for rg_batches in per_rg_results {
            for (sid, batch) in rg_batches {
                per_symbol.entry(sid).or_default().push(batch);
            }
        }

        // Concat and cache each symbol
        let schema = canonical_bars_1m_schema();
        let mut cache = self
            .cache
            .write()
            .map_err(|_| BtError::Data("mega cache lock poisoned".to_string()))?;

        let mut total_rows = 0usize;
        for symbol in symbols {
            let cache_key: CacheKey = (*symbol, range.start.0, range.end.0);
            if let Some(batches) = per_symbol.remove(&symbol.0) {
                let merged = concat_batches(schema.clone(), &batches)?;
                total_rows += merged.num_rows();
                cache.insert(cache_key, merged);
            } else {
                cache.insert(cache_key, RecordBatch::new_empty(schema.clone()));
            }
        }

        tracing::info!(
            total_rows,
            symbols = symbols.len(),
            "mega store: preload complete"
        );

        Ok(())
    }
}

impl DataStore for MegaDataStore {
    fn load_bars(
        &self,
        symbol: SymbolId,
        range: TimeRange,
        _version: Option<&str>,
    ) -> Result<RecordBatch> {
        let cache_key: CacheKey = (symbol, range.start.0, range.end.0);

        // Check cache
        {
            let cache = self
                .cache
                .read()
                .map_err(|_| BtError::Data("mega cache lock poisoned".to_string()))?;
            if let Some(batch) = cache.get(&cache_key) {
                return Ok(batch.clone());
            }

            // Superset check
            for ((sym, cstart, cend), batch) in cache.iter() {
                if *sym == symbol && *cstart <= range.start.0 && *cend >= range.end.0 {
                    let filtered = filter_batch_by_symbol_and_time(batch, symbol, range)?;
                    drop(cache);
                    let mut wc = self
                        .cache
                        .write()
                        .map_err(|_| BtError::Data("mega cache lock poisoned".to_string()))?;
                    wc.insert(cache_key, filtered.clone());
                    return Ok(filtered);
                }
            }
        }

        // Cache miss: read from mmap'd parquet using row group filtering.
        // Single reader for all relevant row groups (footer parsed once).
        let relevant = self.relevant_row_groups(symbol, range);
        tracing::debug!(
            symbol = symbol.0,
            total_rg = self.parquet_metadata.num_row_groups(),
            selected_rg = relevant.len(),
            "mega store: row group pruning"
        );

        let merged = self.read_row_groups(relevant, symbol, range)?;
        tracing::debug!(
            symbol = symbol.0,
            rows = merged.num_rows(),
            "mega store: loaded bars"
        );

        // Cache
        {
            let mut cache = self
                .cache
                .write()
                .map_err(|_| BtError::Data("mega cache lock poisoned".to_string()))?;
            cache.insert(cache_key, merged.clone());
        }

        Ok(merged)
    }

    fn preload(&self, symbols: &[SymbolId], range: TimeRange) -> Result<()> {
        self.preload(symbols, range)
    }

    fn load_features(
        &self,
        _feature_set: &str,
        _symbol: SymbolId,
        _range: TimeRange,
        _columns: Option<&[&str]>,
    ) -> Result<RecordBatch> {
        Err(BtError::Data(
            "MegaDataStore does not support load_features — use ParquetDataStore".to_string(),
        ))
    }

    fn load_exo(&self, _name: &str, _range: TimeRange) -> Result<RecordBatch> {
        Err(BtError::Data(
            "MegaDataStore does not support load_exo — use ParquetDataStore".to_string(),
        ))
    }

    fn list_versions(&self, dataset: &str) -> Result<Vec<DatasetVersion>> {
        self.metadata.list_versions(dataset)
    }

    fn active_version(&self, dataset: &str) -> Result<DatasetVersion> {
        self.metadata.active_version(dataset)
    }

    fn write_bars(
        &self,
        _symbol: SymbolId,
        _batch: &RecordBatch,
        _opts: WriteOptions,
    ) -> Result<DatasetVersion> {
        Err(BtError::Data(
            "MegaDataStore is read-only".to_string(),
        ))
    }
}

/// Filter a batch to rows matching a specific symbol_id AND time range.
fn filter_batch_by_symbol_and_time(
    batch: &RecordBatch,
    symbol: SymbolId,
    range: TimeRange,
) -> Result<RecordBatch> {
    let ts_col = batch
        .column_by_name("timestamp")
        .and_then(|c| c.as_any().downcast_ref::<TimestampNanosecondArray>())
        .ok_or_else(|| BtError::Data("missing/invalid timestamp column".to_string()))?;

    let sym_col = batch
        .column_by_name("symbol_id")
        .and_then(|c| c.as_any().downcast_ref::<UInt32Array>())
        .ok_or_else(|| BtError::Data("missing/invalid symbol_id column".to_string()))?;

    let sid = symbol.0;
    let start = range.start.0;
    let end = range.end.0;

    let mask = BooleanArray::from(
        ts_col
            .values()
            .iter()
            .zip(sym_col.values().iter())
            .map(|(&ts, &s)| s == sid && ts >= start && ts < end)
            .collect::<Vec<bool>>(),
    );

    Ok(filter_record_batch(batch, &mask)?)
}
