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

use arrow::array::{Array, ArrayRef, BooleanArray, TimestampNanosecondArray, UInt32Array};
use arrow::compute::{filter_record_batch, take};
use arrow::ipc::reader::FileReader as IpcFileReader;
use arrow::record_batch::RecordBatch;
use bt_kernel::{
    concat_batches, BtError, Result, SymbolId, TimeRange, Timestamp, NANOS_PER_SECOND,
};
use chrono::Datelike;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::arrow::ArrowWriter;
use parquet::basic::Compression;
use parquet::file::properties::WriterProperties;
use rayon::prelude::*;
use uuid::Uuid;

use crate::metadata::SqliteMetadataStore;
use crate::partition::PartitionScheme;
use crate::schema::{canonical_bars_1m_schema, schema_hash, validate_bars_1m_batch};
use crate::version::DatasetVersion;
use crate::DefaultPartitionScheme;

#[derive(Clone, Debug)]
pub struct WriteOptions {
    pub dataset: String,
    pub source_hash: String,
    pub metadata: serde_json::Value,
    pub set_active: bool,
    pub overwrite_existing: bool,
}

impl Default for WriteOptions {
    fn default() -> Self {
        Self {
            dataset: "bars_1m".to_string(),
            source_hash: "unknown".to_string(),
            metadata: serde_json::json!({}),
            set_active: true,
            overwrite_existing: true,
        }
    }
}

pub trait DataStore: Send + Sync {
    fn load_bars(
        &self,
        symbol: SymbolId,
        range: TimeRange,
        version: Option<&str>,
    ) -> Result<RecordBatch>;

    fn load_features(
        &self,
        feature_set: &str,
        symbol: SymbolId,
        range: TimeRange,
        columns: Option<&[&str]>,
    ) -> Result<RecordBatch>;

    fn load_exo(&self, name: &str, range: TimeRange) -> Result<RecordBatch>;

    fn list_versions(&self, dataset: &str) -> Result<Vec<DatasetVersion>>;

    fn active_version(&self, dataset: &str) -> Result<DatasetVersion>;

    fn write_bars(
        &self,
        symbol: SymbolId,
        batch: &RecordBatch,
        opts: WriteOptions,
    ) -> Result<DatasetVersion>;
}

type BarCacheKey = (SymbolId, i64, i64, Option<String>);

pub struct ParquetDataStore {
    root: PathBuf,
    partition_scheme: Arc<dyn PartitionScheme>,
    metadata: SqliteMetadataStore,
    bar_cache: RwLock<HashMap<BarCacheKey, RecordBatch>>,
    bar_dataset: String,
}

impl ParquetDataStore {
    pub fn new(root: impl AsRef<Path>, metadata_path: impl AsRef<Path>) -> Result<Self> {
        Self::with_dataset(root, metadata_path, "bars_1m")
    }

    pub fn with_dataset(
        root: impl AsRef<Path>,
        metadata_path: impl AsRef<Path>,
        dataset: &str,
    ) -> Result<Self> {
        tracing::info!(root = %root.as_ref().display(), dataset, "opening parquet data store");
        let root = root.as_ref().to_path_buf();
        std::fs::create_dir_all(&root)?;

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

        Ok(Self {
            partition_scheme: Arc::new(DefaultPartitionScheme::new(&root)),
            root,
            metadata,
            bar_cache: RwLock::new(HashMap::new()),
            bar_dataset: dataset.to_string(),
        })
    }

    pub fn with_partition_scheme(
        root: impl AsRef<Path>,
        metadata: SqliteMetadataStore,
        partition_scheme: Arc<dyn PartitionScheme>,
    ) -> Result<Self> {
        let root = root.as_ref().to_path_buf();
        std::fs::create_dir_all(&root)?;

        Ok(Self {
            root,
            partition_scheme,
            metadata,
            bar_cache: RwLock::new(HashMap::new()),
            bar_dataset: "bars_1m".to_string(),
        })
    }

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

impl DataStore for ParquetDataStore {
    #[tracing::instrument(level = "debug", skip(self))]
    fn load_bars(
        &self,
        symbol: SymbolId,
        range: TimeRange,
        version: Option<&str>,
    ) -> Result<RecordBatch> {
        let cache_key: BarCacheKey = (
            symbol,
            range.start.0,
            range.end.0,
            version.map(|v| v.to_string()),
        );

        // Fast path: read lock for exact cache hit
        {
            let cache = self
                .bar_cache
                .read()
                .map_err(|_| BtError::Data("bar cache lock poisoned".to_string()))?;
            if let Some(batch) = cache.get(&cache_key) {
                tracing::debug!(
                    symbol = symbol.0,
                    rows = batch.num_rows(),
                    "cache hit (exact)"
                );
                return Ok(batch.clone());
            }

            // Superset cache: find a cached entry that covers the requested range.
            let version_key = version.map(|v| v.to_string());
            for ((sym, cstart, cend, cver), batch) in cache.iter() {
                if *sym == symbol
                    && *cver == version_key
                    && *cstart <= range.start.0
                    && *cend >= range.end.0
                {
                    let filtered = filter_batch_by_timestamp_range(batch, range)?;
                    tracing::debug!(
                        symbol = symbol.0,
                        rows = filtered.num_rows(),
                        superset_rows = batch.num_rows(),
                        "cache hit (superset slice)"
                    );
                    // Cache the sub-range for next exact hit.
                    drop(cache);
                    let mut wc = self
                        .bar_cache
                        .write()
                        .map_err(|_| BtError::Data("bar cache lock poisoned".to_string()))?;
                    wc.insert(cache_key, filtered.clone());
                    return Ok(filtered);
                }
            }
        }

        // Cache miss: load from disk
        let dataset = &self.bar_dataset;

        if let Some(version_id) = version {
            let exists = self.metadata.version_exists(dataset, version_id)?;
            if !exists {
                return Err(BtError::Data(format!(
                    "requested version '{version_id}' does not exist for dataset '{dataset}'"
                )));
            }
        }

        let partitions = self
            .partition_scheme
            .list_partitions(dataset, symbol, range)?;

        let n_partitions = partitions.len();
        let nested: Vec<Vec<RecordBatch>> = partitions
            .par_iter()
            .enumerate()
            .map(|(idx, partition)| {
                let is_arrow_ipc = partition.extension().is_some_and(|ext| ext == "arrow");

                let file = File::open(partition)?;
                let mut batches = Vec::new();
                let is_boundary = n_partitions <= 2 || idx == 0 || idx == n_partitions - 1;

                if is_arrow_ipc {
                    let reader = IpcFileReader::try_new(file, None)?;
                    for batch_result in reader {
                        let batch = batch_result?;
                        if is_boundary {
                            let filtered = filter_batch_by_timestamp_range(&batch, range)?;
                            if filtered.num_rows() > 0 {
                                batches.push(filtered);
                            }
                        } else {
                            batches.push(batch);
                        }
                    }
                } else {
                    let reader = ParquetRecordBatchReaderBuilder::try_new(file)?
                        .with_batch_size(131_072)
                        .build()?;
                    for batch_result in reader {
                        let batch = batch_result?;
                        if is_boundary {
                            let filtered = filter_batch_by_timestamp_range(&batch, range)?;
                            if filtered.num_rows() > 0 {
                                batches.push(filtered);
                            }
                        } else {
                            batches.push(batch);
                        }
                    }
                }
                Ok(batches)
            })
            .collect::<Result<_>>()?;

        let filtered_batches: Vec<RecordBatch> = nested.into_iter().flatten().collect();

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

        let merged = concat_batches(canonical_bars_1m_schema(), &filtered_batches)?;
        tracing::debug!(
            symbol = symbol.0,
            partitions = n_partitions,
            rows = merged.num_rows(),
            "loaded bars"
        );

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

        Ok(merged)
    }

    fn load_features(
        &self,
        feature_set: &str,
        symbol: SymbolId,
        range: TimeRange,
        columns: Option<&[&str]>,
    ) -> Result<RecordBatch> {
        let dataset = format!("features/{feature_set}");
        let partitions = self
            .partition_scheme
            .list_partitions(&dataset, symbol, range)?;

        if partitions.is_empty() {
            return Err(BtError::Data(format!(
                "no feature partitions found for '{feature_set}' symbol {:?} in range",
                symbol
            )));
        }

        let batches = load_parquet_partitions(&partitions, columns)?;
        if batches.is_empty() {
            return Err(BtError::Data(format!(
                "feature partitions for '{feature_set}' symbol {:?} are empty",
                symbol
            )));
        }

        let schema = batches[0].schema();
        let merged = concat_batches(schema, &batches)?;

        // Filter by time range using timestamp column.
        filter_batch_by_timestamp_range(&merged, range)
    }

    fn load_exo(&self, name: &str, range: TimeRange) -> Result<RecordBatch> {
        let exo_dir = self.root.join("exo").join(name);
        if !exo_dir.is_dir() {
            return Err(BtError::Data(format!(
                "exo data directory not found: {}",
                exo_dir.display()
            )));
        }

        let partitions = list_exo_partitions(&exo_dir, range)?;
        if partitions.is_empty() {
            return Err(BtError::Data(format!(
                "no exo partitions found for '{name}' in range"
            )));
        }

        let batches = load_parquet_partitions(&partitions, None)?;
        if batches.is_empty() {
            return Err(BtError::Data(format!(
                "exo partitions for '{name}' are empty"
            )));
        }

        let schema = batches[0].schema();
        let merged = concat_batches(schema, &batches)?;

        filter_batch_by_timestamp_range(&merged, range)
    }

    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)
    }

    #[tracing::instrument(level = "debug", skip(self, batch))]
    fn write_bars(
        &self,
        symbol: SymbolId,
        batch: &RecordBatch,
        opts: WriteOptions,
    ) -> Result<DatasetVersion> {
        validate_bars_1m_batch(batch)?;
        validate_symbol_column(batch, symbol)?;

        let partitions =
            partition_row_indices(self.partition_scheme.as_ref(), &opts.dataset, symbol, batch)?;

        for (partition, indices) in partitions {
            if partition.exists() && !opts.overwrite_existing {
                return Err(BtError::Data(format!(
                    "partition already exists and overwrite is disabled: {}",
                    partition.display()
                )));
            }

            if let Some(parent) = partition.parent() {
                std::fs::create_dir_all(parent)?;
            }

            let row_ids = UInt32Array::from(indices);
            let partition_batch = take_rows(batch, &row_ids)?;
            write_batch_to_parquet(&partition, &partition_batch)?;
        }

        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!(rows = batch.num_rows(), "wrote bars");
        Ok(version)
    }
}

fn timestamp_column(batch: &RecordBatch) -> Result<&TimestampNanosecondArray> {
    let column = batch
        .column_by_name("timestamp")
        .ok_or_else(|| BtError::Data("missing timestamp column".to_string()))?;

    column
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .ok_or_else(|| BtError::Data("timestamp column has invalid type".to_string()))
}

fn symbol_column(batch: &RecordBatch) -> Result<&UInt32Array> {
    let column = batch
        .column_by_name("symbol_id")
        .ok_or_else(|| BtError::Data("missing symbol_id column".to_string()))?;

    column
        .as_any()
        .downcast_ref::<UInt32Array>()
        .ok_or_else(|| BtError::Data("symbol_id column has invalid type".to_string()))
}

fn validate_symbol_column(batch: &RecordBatch, symbol: SymbolId) -> Result<()> {
    let symbols = symbol_column(batch)?;

    for idx in 0..symbols.len() {
        if symbols.is_null(idx) {
            return Err(BtError::Data(format!("symbol_id at row {idx} is null")));
        }

        if symbols.value(idx) != symbol.0 {
            return Err(BtError::Data(format!(
                "symbol mismatch at row {idx}: expected {}, got {}",
                symbol.0,
                symbols.value(idx)
            )));
        }
    }

    Ok(())
}

fn partition_row_indices(
    scheme: &dyn PartitionScheme,
    dataset: &str,
    symbol: SymbolId,
    batch: &RecordBatch,
) -> Result<BTreeMap<PathBuf, Vec<u32>>> {
    let timestamps = timestamp_column(batch)?;

    let mut groups: BTreeMap<PathBuf, Vec<u32>> = BTreeMap::new();
    for idx in 0..timestamps.len() {
        if timestamps.is_null(idx) {
            return Err(BtError::Data(format!("timestamp at row {idx} is null")));
        }

        let ts = Timestamp(timestamps.value(idx));
        let path = scheme.partition_path(dataset, symbol, ts);
        groups.entry(path).or_default().push(idx as u32);
    }

    Ok(groups)
}

fn take_rows(batch: &RecordBatch, row_ids: &UInt32Array) -> Result<RecordBatch> {
    let mut arrays: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
    for column in batch.columns() {
        arrays.push(take(column.as_ref(), row_ids, None)?);
    }

    Ok(RecordBatch::try_new(batch.schema(), arrays)?)
}

fn write_batch_to_parquet(path: &Path, batch: &RecordBatch) -> Result<()> {
    let file = File::create(path)?;

    let props = WriterProperties::builder()
        .set_compression(Compression::ZSTD(Default::default()))
        .build();

    let mut writer = ArrowWriter::try_new(file, batch.schema(), Some(props))?;
    writer.write(batch)?;
    writer.close()?;

    Ok(())
}

fn compute_time_range(batch: &RecordBatch) -> Result<TimeRange> {
    let timestamps = timestamp_column(batch)?;
    if timestamps.is_empty() {
        return Err(BtError::Data(
            "cannot compute dataset version time range from empty batch".to_string(),
        ));
    }

    let mut min_value = i64::MAX;
    let mut max_value = i64::MIN;

    for idx in 0..timestamps.len() {
        if timestamps.is_null(idx) {
            return Err(BtError::Data(format!("timestamp at row {idx} is null")));
        }

        let ts = timestamps.value(idx);
        min_value = min_value.min(ts);
        max_value = max_value.max(ts);
    }

    let end = max_value
        .checked_add(NANOS_PER_SECOND)
        .ok_or_else(|| BtError::TimeRange("time range end overflow".to_string()))?;

    TimeRange::new(Timestamp(min_value), Timestamp(end))
}

fn utc_now_timestamp() -> Result<Timestamp> {
    let now = chrono::Utc::now();
    let nanos = now
        .timestamp()
        .checked_mul(NANOS_PER_SECOND)
        .and_then(|v| v.checked_add(i64::from(now.timestamp_subsec_nanos())))
        .ok_or_else(|| BtError::Data("timestamp overflow".to_string()))?;

    Ok(Timestamp(nanos))
}

// ---------------------------------------------------------------------------
// Helpers for load_features and load_exo
// ---------------------------------------------------------------------------

/// Read parquet files, optionally projecting to specific columns.
/// Always includes "timestamp" if column projection is specified.
fn load_parquet_partitions(
    paths: &[PathBuf],
    columns: Option<&[&str]>,
) -> Result<Vec<RecordBatch>> {
    let mut all_batches = Vec::new();

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

        let builder = ParquetRecordBatchReaderBuilder::try_new(file)?.with_batch_size(131_072);

        // Column projection: select only the requested columns.
        let reader = if let Some(cols) = columns {
            let parquet_schema = builder.parquet_schema();
            let mut projection_indices = Vec::new();
            for (idx, field) in parquet_schema.columns().iter().enumerate() {
                let name = field.name();
                if name == "timestamp" || cols.contains(&name) {
                    projection_indices.push(idx);
                }
            }
            let mask = parquet::arrow::ProjectionMask::leaves(parquet_schema, projection_indices);
            builder.with_projection(mask).build()?
        } else {
            builder.build()?
        };
        for batch_result in reader {
            let batch = batch_result?;
            if batch.num_rows() > 0 {
                all_batches.push(batch);
            }
        }
    }

    Ok(all_batches)
}

/// Filter a RecordBatch to only include rows within the given time range.
fn filter_batch_by_timestamp_range(batch: &RecordBatch, range: TimeRange) -> Result<RecordBatch> {
    let ts_col = timestamp_column(batch)?;
    let start = range.start.0;
    let end = range.end.0;

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

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

/// List exo parquet files using monthly partitioning: `{exo_dir}/{year}/{month}.parquet`
fn list_exo_partitions(exo_dir: &Path, range: TimeRange) -> Result<Vec<PathBuf>> {
    use crate::partition::timestamp_to_datetime;

    let start_dt = timestamp_to_datetime(range.start)?;
    let last_nanos = range.end.as_nanos().checked_sub(1).ok_or_else(|| {
        BtError::TimeRange("range end underflow while listing exo partitions".to_string())
    })?;
    let end_dt = timestamp_to_datetime(Timestamp(last_nanos))?;

    let start_ym = (start_dt.year(), start_dt.month());
    let end_ym = (end_dt.year(), end_dt.month());

    let mut paths = Vec::new();

    let year_entries = match std::fs::read_dir(exo_dir) {
        Ok(entries) => entries,
        Err(_) => return Ok(Vec::new()),
    };

    for year_entry in year_entries.filter_map(|e| e.ok()) {
        let Some(year) = year_entry
            .file_name()
            .to_str()
            .and_then(|s| s.parse::<i32>().ok())
        else {
            continue;
        };
        if year < start_ym.0 || year > end_ym.0 {
            continue;
        }

        let Ok(month_entries) = std::fs::read_dir(year_entry.path()) else {
            continue;
        };

        for month_entry in month_entries.filter_map(|e| e.ok()) {
            let name = month_entry.file_name();
            let Some(month) = name
                .to_str()
                .and_then(|s| s.strip_suffix(".parquet"))
                .and_then(|s| s.parse::<u32>().ok())
            else {
                continue;
            };

            let ym = (year, month);
            if ym >= start_ym && ym <= end_ym {
                paths.push(month_entry.path());
            }
        }
    }

    paths.sort();
    Ok(paths)
}
