//! DataStore backed by per-symbol Arrow IPC files with multi-resolution support.
//!
//! Directory layout:
//!   {root}/bars_1m/      ← 1m files: {symbol_id}.arrow
//!   {root}/bars_1h/      ← 1h files: {symbol_id}.arrow
//!
//! The store auto-selects the best resolution ≤ requested bar_interval.

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

use arrow::array::{BooleanArray, TimestampNanosecondArray};
use arrow::compute::{filter_record_batch, sort_to_indices, take};
use arrow::ipc::reader::FileReader as IpcFileReader;
use arrow::ipc::writer::FileWriter as IpcFileWriter;
use arrow::record_batch::RecordBatch;
use bt_kernel::{concat_batches, BtError, Result, SymbolId, TimeRange};
use uuid::Uuid;
use rayon::prelude::*;

use crate::metadata::SqliteMetadataStore;
use crate::schema::{canonical_bars_1m_schema, schema_hash, validate_bars_1m_batch};
use crate::store::{compute_time_range, utc_now_timestamp, DataStore};
use crate::version::DatasetVersion;
use crate::WriteOptions;

type CacheKey = (SymbolId, i64, i64, i64); // symbol, start, end, resolution_ns

/// Available resolutions, sorted finest → coarsest.
const RESOLUTIONS: &[(i64, &str)] = &[
    (60_000_000_000, "bars_1m"),    // 1m
    (3_600_000_000_000, "bars_1h"), // 1h
];

/// Read-write DataStore backed by per-symbol Arrow IPC files with multi-resolution support.
///
/// Auto-selects the coarsest resolution that is ≤ the requested bar_interval,
/// minimizing data loaded and alignment work.
pub struct ArrowIpcDataStore {
    root: PathBuf,
    metadata: SqliteMetadataStore,
    cache: RwLock<HashMap<CacheKey, RecordBatch>>,
}

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

        if !root.is_dir() {
            return Err(BtError::Data(format!(
                "Arrow IPC directory not found: {}",
                root.display()
            )));
        }

        // Log available resolutions
        for &(ns, subdir) in RESOLUTIONS {
            let dir = root.join(subdir);
            let label = match ns {
                60_000_000_000 => "1m",
                3_600_000_000_000 => "1h",
                _ => "?",
            };
            let exists = dir.is_dir()
                && std::fs::read_dir(&dir)
                    .map(|mut d| d.any(|e| e.map(|e| e.path().extension().is_some_and(|ext| ext == "arrow")).unwrap_or(false)))
                    .unwrap_or(false);
            if exists {
                tracing::info!(resolution = label, dir = %dir.display(), "resolution available");
            }
        }

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

        Ok(Self {
            root,
            metadata,
            cache: RwLock::new(HashMap::new()),
        })
    }

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

    /// Read all batches from an Arrow IPC file (no time filtering).
    fn read_symbol_from_file(path: &Path) -> Result<RecordBatch> {
        let file = File::open(path)
            .map_err(|e| BtError::Data(format!("opening {}: {e}", path.display())))?;
        let reader = IpcFileReader::try_new(file, None)
            .map_err(|e| BtError::Data(format!("reading IPC {}: {e}", path.display())))?;

        let mut batches = Vec::new();
        for batch_result in reader {
            let batch = batch_result
                .map_err(|e| BtError::Data(format!("reading IPC batch: {e}")))?;
            batches.push(batch);
        }

        if batches.is_empty() {
            return Ok(RecordBatch::new_empty(canonical_bars_1m_schema()));
        }
        let schema = batches[0].schema();
        concat_batches(schema, &batches)
    }

    /// Sort a RecordBatch by timestamp and remove duplicate timestamps.
    fn sort_and_dedup(batch: &RecordBatch) -> Result<RecordBatch> {
        // Sort by timestamp.
        let ts_col = batch
            .column_by_name("timestamp")
            .ok_or_else(|| BtError::Data("missing timestamp column".into()))?;

        let sort_opts = arrow::compute::SortOptions {
            descending: false,
            nulls_first: false,
        };
        let indices = sort_to_indices(ts_col, Some(sort_opts), None)?;

        let sorted_columns: Vec<_> = batch
            .columns()
            .iter()
            .map(|col| take(col, &indices, None).map_err(BtError::from))
            .collect::<Result<_>>()?;

        let sorted = RecordBatch::try_new(batch.schema(), sorted_columns)
            .map_err(|e| BtError::Data(format!("rebuilding sorted batch: {e}")))?;

        // Dedup: keep first occurrence of each timestamp.
        let ts_arr = sorted
            .column_by_name("timestamp")
            .unwrap()
            .as_any()
            .downcast_ref::<TimestampNanosecondArray>()
            .unwrap();

        let mut keep = vec![true; ts_arr.len()];
        for (i, k) in keep.iter_mut().enumerate().skip(1) {
            if ts_arr.value(i) == ts_arr.value(i - 1) {
                *k = false;
            }
        }

        let mask = BooleanArray::from(keep);
        let deduped = filter_record_batch(&sorted, &mask)?;
        Ok(deduped)
    }

    /// Pick the best (coarsest) resolution that is ≤ bar_interval_ns.
    /// Returns (resolution_ns, directory_path).
    fn best_resolution(&self, bar_interval_ns: i64) -> (i64, PathBuf) {
        // Walk from coarsest to finest, pick the first one that:
        // 1. Is ≤ bar_interval_ns
        // 2. Has .arrow files in its directory
        for &(res_ns, subdir) in RESOLUTIONS.iter().rev() {
            if res_ns > bar_interval_ns {
                continue;
            }
            let dir = self.root.join(subdir);
            if dir.is_dir() {
                // Check at least one .arrow file exists
                let has_files = std::fs::read_dir(&dir)
                    .map(|mut d| d.any(|e| {
                        e.map(|e| e.path().extension().is_some_and(|ext| ext == "arrow"))
                            .unwrap_or(false)
                    }))
                    .unwrap_or(false);
                if has_files {
                    return (res_ns, dir);
                }
            }
        }

        // Fallback: finest (1m)
        (RESOLUTIONS[0].0, self.root.join(RESOLUTIONS[0].1))
    }

    /// Read a symbol's Arrow IPC file from a specific directory, filtering to time range.
    fn read_symbol_from_dir(
        dir: &Path,
        symbol: SymbolId,
        range: TimeRange,
    ) -> Result<RecordBatch> {
        let path = dir.join(format!("{}.arrow", symbol.0));
        if !path.exists() {
            return Ok(RecordBatch::new_empty(canonical_bars_1m_schema()));
        }

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

        let reader = IpcFileReader::try_new(file, None)
            .map_err(|e| BtError::Data(format!("reading IPC {}: {e}", path.display())))?;

        let start = range.start.0;
        let end = range.end.0;
        let mut all_batches = Vec::new();

        for batch_result in reader {
            let batch = batch_result
                .map_err(|e| BtError::Data(format!("reading IPC batch: {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()))?;

            // Fast path: check if entire batch is in range
            let first_ts = ts_col.value(0);
            let last_ts = ts_col.value(ts_col.len() - 1);

            if first_ts >= start && last_ts < end {
                all_batches.push(batch);
            } else if last_ts < start || first_ts >= end {
                continue;
            } else {
                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 {
                    all_batches.push(filtered);
                }
            }
        }

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

        let schema = all_batches[0].schema();
        concat_batches(schema, &all_batches)
    }
}

impl DataStore for ArrowIpcDataStore {
    fn load_bars(
        &self,
        symbol: SymbolId,
        range: TimeRange,
        _version: Option<&str>,
    ) -> Result<RecordBatch> {
        self.load_bars_resolution(symbol, range, 60_000_000_000) // default: 1m
    }

    fn load_bars_resolution(
        &self,
        symbol: SymbolId,
        range: TimeRange,
        bar_interval_ns: i64,
    ) -> Result<RecordBatch> {
        let (res_ns, dir) = self.best_resolution(bar_interval_ns);
        let cache_key: CacheKey = (symbol, range.start.0, range.end.0, res_ns);

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

        let batch = Self::read_symbol_from_dir(&dir, symbol, range)?;

        let res_label = match res_ns {
            60_000_000_000 => "1m",
            3_600_000_000_000 => "1h",
            _ => "?",
        };
        tracing::debug!(
            symbol = symbol.0,
            resolution = res_label,
            rows = batch.num_rows(),
            "Arrow IPC: loaded bars"
        );

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

        Ok(batch)
    }

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

    fn preload_resolution(
        &self,
        symbols: &[SymbolId],
        range: TimeRange,
        bar_interval_ns: i64,
    ) -> Result<()> {
        let (res_ns, dir) = self.best_resolution(bar_interval_ns);

        let res_label = match res_ns {
            60_000_000_000 => "1m",
            3_600_000_000_000 => "1h",
            _ => "?",
        };
        tracing::info!(
            symbols = symbols.len(),
            resolution = res_label,
            "Arrow IPC store: parallel preload"
        );

        let results: Vec<(SymbolId, RecordBatch)> = symbols
            .par_iter()
            .map(|&symbol| {
                let batch = Self::read_symbol_from_dir(&dir, symbol, range)?;
                Ok((symbol, batch))
            })
            .collect::<Result<_>>()?;

        let mut cache = self
            .cache
            .write()
            .map_err(|_| BtError::Data("cache lock poisoned".to_string()))?;

        let mut total_rows = 0usize;
        for (symbol, batch) in results {
            total_rows += batch.num_rows();
            let cache_key: CacheKey = (symbol, range.start.0, range.end.0, res_ns);
            cache.insert(cache_key, batch);
        }

        tracing::info!(total_rows, "Arrow IPC store: preload complete");
        Ok(())
    }

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

    fn load_exo(&self, _name: &str, _range: TimeRange) -> Result<RecordBatch> {
        Err(BtError::Data(
            "ArrowIpcDataStore does not support load_exo".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> {
        validate_bars_1m_batch(batch)?;

        // Merge with existing data if present.
        let dir = self.root.join("bars_1m");
        std::fs::create_dir_all(&dir)?;
        let path = dir.join(format!("{}.arrow", symbol.0));

        let merged = if path.exists() {
            // Read existing data.
            let existing = Self::read_symbol_from_file(&path)?;
            if existing.num_rows() > 0 {
                // Concat existing + new, sort by timestamp, deduplicate.
                let combined = concat_batches(batch.schema(), &[existing, batch.clone()])?;
                Self::sort_and_dedup(&combined)?
            } else {
                batch.clone()
            }
        } else {
            batch.clone()
        };

        let file = std::fs::File::create(&path)
            .map_err(|e| BtError::Data(format!("creating {}: {e}", path.display())))?;

        let mut writer = IpcFileWriter::try_new(file, merged.schema().as_ref())
            .map_err(|e| BtError::Data(format!("IPC writer init: {e}")))?;
        writer
            .write(&merged)
            .map_err(|e| BtError::Data(format!("IPC write: {e}")))?;
        writer
            .finish()
            .map_err(|e| BtError::Data(format!("IPC finish: {e}")))?;
        let batch = &merged;

        // Invalidate cache for this symbol
        if let Ok(mut cache) = self.cache.write() {
            cache.retain(|k, _| k.0 != symbol);
        }

        let time_range = compute_time_range(batch)?;

        let version = DatasetVersion {
            id: Uuid::new_v4().to_string(),
            dataset: opts.dataset.clone(),
            created_at: utc_now_timestamp()?,
            schema_hash: schema_hash(batch.schema().as_ref()),
            row_count: batch.num_rows() as u64,
            time_range,
            symbols: vec![symbol],
            source_hash: opts.source_hash,
            metadata: opts.metadata,
        };

        self.metadata.register_dataset_version(
            &version.dataset,
            &version,
            &self.root,
            opts.set_active,
        )?;

        tracing::debug!(symbol = symbol.0, rows = batch.num_rows(), path = %path.display(), "wrote Arrow IPC bars");
        Ok(version)
    }
}
