//! Feature Cache for DQN Training
//!
//! This module provides a disk-based caching system for computed feature vectors,
//! eliminating redundant MBP-10 OFI calculations across hyperopt trials.
//!
//! ## Cache Architecture
//!
//! - **Storage Format**: Apache Parquet (columnar, Snappy compression)
//! - **Cache Key**: SHA256 hash of (parquet path, mbp10 dir, warmup, extraction.rs)
//! - **Metadata**: JSON index mapping cache keys to file locations
//! - **Expected Savings**: 8-25 minutes per 50-trial hyperopt campaign
//!
//! ## Usage
//!
//! ```rust
//! use ml::feature_cache::{calculate_cache_key, load_features_from_cache};
//!
//! let cache_key = calculate_cache_key(
//!     Path::new("test_data/ES_FUT_180d.parquet"),
//!     Path::new("test_data/mbp10"),
//!     50, // warmup period
//! )?;
//!
//! if let Some(features) = load_features_from_cache(cache_dir, &cache_key).await? {
//!     // Use cached features
//! } else {
//!     // Compute features and save to cache
//! }
//! ```

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::SystemTime;
use tracing::{debug, info, warn};


/// Cache metadata for a single cached feature set
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheMetadata {
    /// Cache format version
    pub version: String,
    /// SHA256 cache key
    pub cache_key: String,
    /// Parquet file path
    pub parquet_path: String,
    /// Parquet file modification time (Unix timestamp)
    pub parquet_mtime: u64,
    /// MBP-10 directory path
    pub mbp10_dir: String,
    /// Number of .dbn files in MBP-10 directory
    pub mbp10_file_count: usize,
    /// MBP-10 file modification times (Unix timestamps)
    pub mbp10_mtimes: Vec<u64>,
    /// Number of features per sample (51)
    pub feature_count: usize,
    /// Number of bars cached
    pub bar_count: usize,
    /// Warmup period used
    pub warmup_period: usize,
    /// Cache creation timestamp (ISO 8601)
    pub created_at: String,
}

/// Calculate cache key from input parameters
///
/// Generates a SHA256 hash that uniquely identifies a cached feature set.
/// The hash includes:
/// 1. Parquet file path + modification time
/// 2. MBP-10 directory + all .dbn files + modification times
/// 3. Feature extraction code (extraction.rs hash)
/// 4. Warmup period constant
///
/// # Arguments
///
/// * `parquet_path` - Path to OHLCV parquet file
/// * `mbp10_dir` - Path to MBP-10 snapshot directory
/// * `warmup_period` - Warmup period for feature extraction (typically 50)
///
/// # Returns
///
/// Hex-encoded SHA256 hash (64 characters)
pub fn calculate_cache_key(
    parquet_path: &Path,
    mbp10_dir: &Path,
    warmup_period: usize,
) -> Result<String> {
    let mut hasher = Sha256::new();

    // 1. Parquet file path + mtime
    hasher.update(parquet_path.to_string_lossy().as_bytes());
    let parquet_mtime = parquet_path
        .metadata()
        .context("Failed to read parquet metadata")?
        .modified()
        .context("Failed to get parquet mtime")?
        .duration_since(SystemTime::UNIX_EPOCH)
        .context("Failed to convert parquet mtime")?
        .as_secs();
    hasher.update(&parquet_mtime.to_le_bytes());

    // 2. MBP-10 directory + all .dbn files + mtimes
    hasher.update(mbp10_dir.to_string_lossy().as_bytes());
    let mut dbn_files: Vec<_> = std::fs::read_dir(mbp10_dir)
        .context("Failed to read mbp10 directory")?
        .filter_map(|e| e.ok())
        .filter(|e| {
            e.path()
                .extension()
                .and_then(|s| s.to_str())
                == Some("dbn")
        })
        .collect();
    dbn_files.sort_by_key(|e| e.path());

    for entry in &dbn_files {
        hasher.update(entry.path().to_string_lossy().as_bytes());
        let mtime = entry
            .metadata()
            .context("Failed to read dbn metadata")?
            .modified()
            .context("Failed to get dbn mtime")?
            .duration_since(SystemTime::UNIX_EPOCH)
            .context("Failed to convert dbn mtime")?
            .as_secs();
        hasher.update(&mtime.to_le_bytes());
    }

    // 3. Feature extraction version (hash of extraction.rs source)
    let extraction_src_path = PathBuf::from("ml/src/features/extraction.rs");
    if extraction_src_path.exists() {
        let extraction_src = std::fs::read_to_string(&extraction_src_path)
            .context("Failed to read extraction.rs")?;
        hasher.update(extraction_src.as_bytes());
    } else {
        warn!(
            "extraction.rs not found at {:?}, cache key may be incomplete",
            extraction_src_path
        );
    }

    // 4. Warmup period
    hasher.update(&warmup_period.to_le_bytes());

    Ok(format!("{:x}", hasher.finalize()))
}

/// Load features from cache
///
/// Reads cached feature vectors from Parquet file if cache hit.
/// Returns None if cache miss or validation fails.
///
/// # Arguments
///
/// * `cache_dir` - Directory containing cache files
/// * `cache_key` - SHA256 cache key from calculate_cache_key()
///
/// # Returns
///
/// * `Ok(Some(features))` - Cache hit, features loaded (Vec<[f64; 42]>)
/// * `Ok(None)` - Cache miss or validation failed
/// * `Err(e)` - I/O or parsing error
pub async fn load_features_from_cache(
    cache_dir: &Path,
    cache_key: &str,
) -> Result<Option<Vec<[f64; 42]>>> {
    use arrow::array::Float64Array;
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    use std::fs::File;

    // Check metadata exists
    let metadata_path = cache_dir.join("cache_metadata.json");
    if !metadata_path.exists() {
        debug!("Cache miss: no metadata file found");
        return Ok(None);
    }

    // Load and verify metadata
    let file = File::open(&metadata_path).context("Failed to open cache metadata")?;
    let metadata_map: HashMap<String, CacheMetadata> =
        serde_json::from_reader(file).context("Failed to parse cache metadata")?;

    let metadata = match metadata_map.get(cache_key) {
        Some(m) => m,
        None => {
            debug!("Cache miss: no cached features for cache key {}", cache_key);
            return Ok(None);
        }
    };

    // Check cache file exists
    let cache_path = cache_dir.join(format!("features_{}.parquet", cache_key));
    if !cache_path.exists() {
        debug!("Cache file not present: {:?}", cache_path);
        return Ok(None);
    }

    info!("🔍 Cache hit! Loading features from {:?}", cache_path);

    // Load features from Parquet
    let file = File::open(&cache_path).context("Failed to open cache file")?;
    let builder = ParquetRecordBatchReaderBuilder::try_new(file)
        .context("Failed to create Parquet reader")?;
    let reader = builder.build().context("Failed to build Parquet reader")?;

    let mut all_features = Vec::new();

    for batch_result in reader {
        let batch = batch_result.context("Failed to read Parquet batch")?;
        let num_rows = batch.num_rows();

        for row_idx in 0..num_rows {
            let mut feature_vec = [0.0; 42];

            for col_idx in 0..42 {
                let column = batch.column(col_idx);
                let array = column
                    .as_any()
                    .downcast_ref::<Float64Array>()
                    .context("Invalid column type in cache")?;
                feature_vec[col_idx] = array.value(row_idx);
            }

            all_features.push(feature_vec);
        }
    }

    info!(
        "✅ Loaded {} feature vectors from cache (estimated savings: ~20.0s)",
        all_features.len()
    );
    info!("   Cache metadata: {} bars × {} features", metadata.bar_count, metadata.feature_count);

    Ok(Some(all_features))
}

/// Save features to cache
///
/// Writes computed feature vectors to Parquet cache file and updates metadata.
///
/// # Arguments
///
/// * `cache_dir` - Directory to store cache files
/// * `cache_key` - SHA256 cache key
/// * `features` - Feature vectors to cache ([f64; 42])
/// * `metadata` - Cache metadata
///
/// # Returns
///
/// * `Ok(())` - Cache saved successfully
/// * `Err(e)` - I/O or serialization error
pub async fn save_features_to_cache(
    cache_dir: &Path,
    cache_key: &str,
    features: &[[f64; 42]],
    metadata: CacheMetadata,
) -> Result<()> {
    use arrow::array::Float64Array;
    use arrow::datatypes::{DataType, Field, Schema};
    use arrow::record_batch::RecordBatch;
    use parquet::arrow::ArrowWriter;
    use parquet::file::properties::WriterProperties;
    use std::fs::File;
    use std::sync::Arc;

    // Create cache directory if it doesn't exist
    std::fs::create_dir_all(cache_dir).context("Failed to create cache directory")?;

    // Define Parquet schema (51 float64 columns)
    let mut fields = Vec::with_capacity(51);
    for i in 0..51 {
        fields.push(Field::new(format!("feature_{}", i), DataType::Float64, false));
    }
    let schema = Arc::new(Schema::new(fields));

    // Create Parquet writer with Snappy compression
    let cache_path = cache_dir.join(format!("features_{}.parquet", cache_key));
    let file = File::create(&cache_path).context("Failed to create cache file")?;
    let props = WriterProperties::builder()
        .set_compression(parquet::basic::Compression::SNAPPY)
        .build();
    let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))
        .context("Failed to create Parquet writer")?;

    // Convert features to Arrow arrays
    let mut columns: Vec<Arc<dyn arrow::array::Array>> = Vec::with_capacity(51);
    for col_idx in 0..51 {
        let col_data: Vec<f64> = features.iter().map(|row| row[col_idx]).collect();
        columns.push(Arc::new(Float64Array::from(col_data)));
    }

    // Write batch to Parquet
    let batch = RecordBatch::try_new(schema, columns).context("Failed to create record batch")?;
    writer.write(&batch).context("Failed to write Parquet batch")?;
    writer.close().context("Failed to close Parquet writer")?;

    info!("💾 Saved {} features to cache: {:?}", features.len(), cache_path);

    // Update metadata index
    let metadata_path = cache_dir.join("cache_metadata.json");
    let mut metadata_map: HashMap<String, CacheMetadata> = if metadata_path.exists() {
        let file = File::open(&metadata_path).context("Failed to open metadata for update")?;
        serde_json::from_reader(file).unwrap_or_default()
    } else {
        HashMap::new()
    };

    metadata_map.insert(cache_key.to_string(), metadata);

    let file = File::create(&metadata_path).context("Failed to create metadata file")?;
    serde_json::to_writer_pretty(file, &metadata_map)
        .context("Failed to write metadata")?;

    info!("✅ Cache metadata updated: {:?}", metadata_path);

    Ok(())
}

// ══════════════════════════════════════════════════════════════════════════════
// DBN Training Data Cache — for smoke test speedup
//
// Caches the full output of `load_training_data()` — the merged
// `Vec<([f64; 42], Vec<f64>)>` — to `/tmp/.foxhunt_feature_cache/` so that
// repeated test runs skip the ~60s DBN decode + feature extraction step.
//
// Cache key: SHA256 of all .dbn / .dbn.zst filenames + sizes + mtimes in the
// data directory (recursive). Any file change invalidates the cache.
//
// Disable with `FOXHUNT_FEATURE_CACHE=0`.
// ══════════════════════════════════════════════════════════════════════════════

/// Calculate a cache key for a DBN data directory.
///
/// Hashes every `.dbn` and `.dbn.zst` file found recursively under `data_dir`
/// by (canonical path, size, mtime). Any change to the directory contents
/// (add / remove / modify) produces a different key.
pub fn calculate_dbn_cache_key(data_dir: &Path) -> Result<String> {
    calculate_dbn_cache_key_full(data_dir, None, None)
}

/// Extended cache key that includes MBP-10 and trades directories.
/// The cache is invalidated when ANY data source changes — OHLCV, MBP-10, or trades.
/// Without this, a cache built without trades would serve stale features
/// even after trades data is added, silently dropping VPIN/Kyle's Lambda enrichment.
pub fn calculate_dbn_cache_key_full(
    data_dir: &Path,
    mbp10_dir: Option<&Path>,
    trades_dir: Option<&Path>,
) -> Result<String> {
    let mut hasher = Sha256::new();
    hasher.update(data_dir.to_string_lossy().as_bytes());

    let mut files: Vec<_> = collect_dbn_files_for_hash(data_dir);
    // Also include MBP-10 and trades files in the hash
    if let Some(dir) = mbp10_dir {
        hasher.update(b"mbp10:");
        hasher.update(dir.to_string_lossy().as_bytes());
        files.extend(collect_dbn_files_for_hash(dir));
    } else {
        hasher.update(b"mbp10:none");
    }
    if let Some(dir) = trades_dir {
        hasher.update(b"trades:");
        hasher.update(dir.to_string_lossy().as_bytes());
        files.extend(collect_dbn_files_for_hash(dir));
    } else {
        hasher.update(b"trades:none");
    }
    files.sort(); // deterministic ordering

    if files.is_empty() {
        return Err(anyhow::anyhow!(
            "No .dbn or .dbn.zst files found in {:?} — cannot build cache key",
            data_dir
        ));
    }

    for path in &files {
        hasher.update(path.to_string_lossy().as_bytes());
        let meta = path
            .metadata()
            .with_context(|| format!("Failed to stat {:?}", path))?;
        hasher.update(&meta.len().to_le_bytes());
        let mtime = meta
            .modified()
            .context("mtime unavailable")?
            .duration_since(SystemTime::UNIX_EPOCH)
            .context("mtime before epoch")?
            .as_secs();
        hasher.update(&mtime.to_le_bytes());
    }

    Ok(format!("{:x}", hasher.finalize()))
}

/// Collect all .dbn and .dbn.zst paths recursively (for hashing purposes).
fn collect_dbn_files_for_hash(dir: &Path) -> Vec<PathBuf> {
    let mut out = Vec::new();
    if let Ok(entries) = std::fs::read_dir(dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                out.extend(collect_dbn_files_for_hash(&path));
            } else if path.extension().and_then(|s| s.to_str()) == Some("dbn")
                || path.to_string_lossy().ends_with(".dbn.zst")
            {
                out.push(path);
            }
        }
    }
    out
}

/// Return the cache directory, or `None` when caching is disabled.
///
/// Resolution: `$CARGO_TARGET_DIR/.foxhunt_feature_cache` (PVC-persisted on CI),
/// then `/tmp/.foxhunt_feature_cache` (ephemeral fallback).
/// Disabled with `FOXHUNT_FEATURE_CACHE=0`.
fn dbn_cache_dir() -> Option<PathBuf> {
    if std::env::var("FOXHUNT_FEATURE_CACHE")
        .map(|v| v == "0")
        .unwrap_or(false)
    {
        debug!("FOXHUNT_FEATURE_CACHE=0 — DBN feature cache disabled");
        return None;
    }
    let dir = if let Ok(target_dir) = std::env::var("CARGO_TARGET_DIR") {
        PathBuf::from(target_dir).join(".foxhunt_feature_cache")
    } else {
        PathBuf::from("/tmp/.foxhunt_feature_cache")
    };
    Some(dir)
}

/// Try to load a cached training dataset for `data_dir`.
///
/// Returns `Ok(Some((train, val)))` on cache hit, `Ok(None)` on miss.
/// The caller should fall through to full data loading on `Ok(None)`.
///
/// # Binary format
///
/// ```text
/// [u64 le] total sample count N
/// for each sample:
///   [42 × f64 le]  feature vector
///   [u64 le]       target vector length T
///   [T  × f64 le]  target values
/// ```
pub fn load_dbn_training_cache(
    data_dir: &Path,
) -> Result<Option<(Vec<([f64; 42], Vec<f64>)>, Vec<([f64; 42], Vec<f64>)>)>> {
    load_dbn_training_cache_full(data_dir, None, None)
}

/// Load cached features with full data source awareness (OHLCV + MBP-10 + trades).
pub fn load_dbn_training_cache_full(
    data_dir: &Path,
    mbp10_dir: Option<&Path>,
    trades_dir: Option<&Path>,
) -> Result<Option<(Vec<([f64; 42], Vec<f64>)>, Vec<([f64; 42], Vec<f64>)>)>> {
    use std::io::{BufReader, Read};

    let cache_dir = match dbn_cache_dir() {
        Some(d) => d,
        None => return Ok(None),
    };

    let key = match calculate_dbn_cache_key_full(data_dir, mbp10_dir, trades_dir) {
        Ok(k) => k,
        Err(e) => {
            debug!("DBN cache key error (skipping cache): {e}");
            return Ok(None);
        }
    };

    let cache_path = cache_dir.join(format!("dbn_features_{key}.bin"));
    if !cache_path.exists() {
        debug!("DBN cache miss: {:?}", cache_path);
        return Ok(None);
    }

    let start = std::time::Instant::now();
    let file = std::fs::File::open(&cache_path)
        .with_context(|| format!("Failed to open DBN cache {:?}", cache_path))?;
    let mut reader = BufReader::new(file);

    // Read total sample count
    let mut count_buf = [0u8; 8];
    reader
        .read_exact(&mut count_buf)
        .with_context(|| "Failed to read sample count from DBN cache")?;
    let n = u64::from_le_bytes(count_buf) as usize;

    let mut data: Vec<([f64; 42], Vec<f64>)> = Vec::with_capacity(n);
    let mut f64_buf = [0u8; 8];

    for _ in 0..n {
        // Read 42-element feature vector
        let mut features = [0.0_f64; 42];
        for feat in &mut features {
            reader
                .read_exact(&mut f64_buf)
                .with_context(|| "Failed to read feature from DBN cache")?;
            *feat = f64::from_le_bytes(f64_buf);
        }

        // Read target vector
        reader
            .read_exact(&mut count_buf)
            .with_context(|| "Failed to read target len from DBN cache")?;
        let t_len = u64::from_le_bytes(count_buf) as usize;
        let mut targets = vec![0.0_f64; t_len];
        for tgt in &mut targets {
            reader
                .read_exact(&mut f64_buf)
                .with_context(|| "Failed to read target value from DBN cache")?;
            *tgt = f64::from_le_bytes(f64_buf);
        }

        data.push((features, targets));
    }

    let elapsed = start.elapsed().as_secs_f64();
    info!(
        "DBN feature cache hit: loaded {} samples in {:.2}s from {:?}",
        data.len(),
        elapsed,
        cache_path
    );

    // Reproduce the same 80/20 split that `load_training_data` uses
    let split_idx = (data.len() * 80) / 100;
    let train = data[..split_idx].to_vec();
    let val = data[split_idx..].to_vec();

    Ok(Some((train, val)))
}

/// Persist a merged training dataset to disk for future cache hits.
///
/// Merges `train_data` and `val_data` back into a single `Vec` (in order)
/// before writing, so `load_dbn_training_cache` can re-split consistently.
///
/// Silently logs and returns on any error — cache misses are always safe.
///
/// See `load_dbn_training_cache` for the binary format specification.
pub fn save_dbn_training_cache(
    data_dir: &Path,
    train_data: &[([f64; 42], Vec<f64>)],
    val_data: &[([f64; 42], Vec<f64>)],
) {
    save_dbn_training_cache_full(data_dir, None, None, train_data, val_data);
}

/// Save cached features with full data source awareness.
pub fn save_dbn_training_cache_full(
    data_dir: &Path,
    mbp10_dir: Option<&Path>,
    trades_dir: Option<&Path>,
    train_data: &[([f64; 42], Vec<f64>)],
    val_data: &[([f64; 42], Vec<f64>)],
) {
    use std::io::{BufWriter, Write};

    let cache_dir = match dbn_cache_dir() {
        Some(d) => d,
        None => return,
    };

    let key = match calculate_dbn_cache_key_full(data_dir, mbp10_dir, trades_dir) {
        Ok(k) => k,
        Err(e) => {
            warn!("DBN cache key error (skipping save): {e}");
            return;
        }
    };

    if let Err(e) = std::fs::create_dir_all(&cache_dir) {
        warn!("Failed to create DBN cache dir {:?}: {e}", cache_dir);
        return;
    }

    let cache_path = cache_dir.join(format!("dbn_features_{key}.bin"));

    let result: Result<u64> = (|| {
        let file = std::fs::File::create(&cache_path)
            .with_context(|| format!("Failed to create DBN cache {:?}", cache_path))?;
        let mut writer = BufWriter::new(file);

        let total = (train_data.len() + val_data.len()) as u64;
        writer.write_all(&total.to_le_bytes())?;

        for (features, targets) in train_data.iter().chain(val_data.iter()) {
            for &v in features {
                writer.write_all(&v.to_le_bytes())?;
            }
            let t_len = targets.len() as u64;
            writer.write_all(&t_len.to_le_bytes())?;
            for &v in targets {
                writer.write_all(&v.to_le_bytes())?;
            }
        }

        writer.flush()?;
        Ok(total)
    })();

    match result {
        Ok(total) => {
            let size_mb = cache_path
                .metadata()
                .map(|m| m.len() as f64 / 1_048_576.0)
                .unwrap_or(0.0);
            info!(
                "DBN feature cache saved: {} samples ({:.1} MB) to {:?}",
                total, size_mb, cache_path
            );
        }
        Err(e) => warn!("DBN cache write error (non-fatal): {e}"),
    }
}

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

    #[test]
    fn test_cache_key_deterministic() {
        // Cache key should be deterministic for same inputs
        let parquet = PathBuf::from("test_data/ES_FUT_180d.parquet");
        let mbp10 = PathBuf::from("test_data/mbp10");

        if parquet.exists() && mbp10.exists() {
            let key1 = calculate_cache_key(&parquet, &mbp10, 50).unwrap();
            let key2 = calculate_cache_key(&parquet, &mbp10, 50).unwrap();
            assert_eq!(key1, key2, "Cache keys should be deterministic");
            assert_eq!(key1.len(), 64, "SHA256 hash should be 64 hex chars");
        }
    }

    #[test]
    fn test_cache_key_varies_with_inputs() {
        let parquet = PathBuf::from("test_data/ES_FUT_180d.parquet");
        let mbp10 = PathBuf::from("test_data/mbp10");

        if parquet.exists() && mbp10.exists() {
            let key1 = calculate_cache_key(&parquet, &mbp10, 50).unwrap();
            let key2 = calculate_cache_key(&parquet, &mbp10, 100).unwrap();
            assert_ne!(key1, key2, "Cache keys should vary with warmup period");
        }
    }
}
