//! DatasetManager -- orchestrates data download, caching, and preparation.
//!
//! The manager consults the [`CacheManifest`] to find which date ranges are
//! already cached, loads the corresponding DBN files, and splits the result
//! into train / val / test sets according to the [`SplitRatio`] in the
//! supplied [`DatasetSpec`].

use super::cache::{CacheManifest, CachedRange};
use super::{BarSize, DatasetSpec, SplitRatio};
use crate::MLError;
use std::path::{Path, PathBuf};
use tracing::{debug, info, warn};

/// Width of the feature vector produced by the data pipeline.
pub const FEATURE_DIM: usize = 128;

/// Maximum number of feature vectors held in memory during `prepare()`.
/// At 128 × f64 (1 KiB per vector) this caps RAM usage at ~512 MB.
const MAX_FEATURES: usize = 500_000;

// ---------------------------------------------------------------------------
// DataSplit
// ---------------------------------------------------------------------------

/// The result of preparing a dataset: features split into train / val / test.
#[derive(Debug, Clone)]
pub struct DataSplit {
    /// Human-readable name carried over from the [`DatasetSpec`].
    pub name: String,
    /// Training set (chronologically first).
    pub train: Vec<[f64; FEATURE_DIM]>,
    /// Validation set (middle).
    pub val: Vec<[f64; FEATURE_DIM]>,
    /// Test set (chronologically last).
    pub test: Vec<[f64; FEATURE_DIM]>,
    /// Total number of bars before splitting.
    pub total_bars: usize,
    /// Symbols that contributed data to this split.
    pub symbols: Vec<String>,
}

// ---------------------------------------------------------------------------
// DatasetManager
// ---------------------------------------------------------------------------

/// Orchestrates loading cached market data and splitting it for ML training.
pub struct DatasetManager {
    cache_dir: PathBuf,
    manifest: CacheManifest,
}

impl DatasetManager {
    /// Create a new manager rooted at `cache_dir`.
    ///
    /// Loads the existing [`CacheManifest`] from disk (or starts with an empty
    /// one if the directory does not exist yet).
    pub fn new(cache_dir: PathBuf) -> Result<Self, MLError> {
        let manifest = CacheManifest::load(&cache_dir)?;
        Ok(Self { cache_dir, manifest })
    }

    /// Path to the cache directory.
    pub fn cache_dir(&self) -> &Path {
        &self.cache_dir
    }

    /// Borrow the current cache manifest.
    pub fn manifest(&self) -> &CacheManifest {
        &self.manifest
    }

    /// Load and split data according to `spec`.
    ///
    /// For each symbol in the spec the manager:
    /// 1. Checks the cache manifest for gaps (logged as warnings).
    /// 2. Loads all cached Parquet files that overlap the requested range.
    /// 3. Concatenates features and splits chronologically.
    ///
    /// Returns an error if no data could be loaded for any symbol.
    pub fn prepare(&self, spec: &DatasetSpec) -> Result<DataSplit, MLError> {
        info!(
            "Preparing dataset '{}': {} symbols",
            spec.name,
            spec.symbols.len()
        );

        let all_features: Vec<[f64; FEATURE_DIM]> = Vec::new();
        let mut loaded_symbols = Vec::new();
        let (req_start, req_end) = spec.date_range.resolve();

        for sym_spec in &spec.symbols {
            // Report gaps so the caller knows what is missing.
            let gaps = self
                .manifest
                .find_gaps(&sym_spec.symbol, req_start, req_end);
            if !gaps.is_empty() {
                warn!(
                    "Symbol {} has {} data gap(s)",
                    sym_spec.symbol,
                    gaps.len()
                );
                for (gs, ge) in &gaps {
                    warn!("  Gap: {} to {}", gs, ge);
                }
            }

            let files = self
                .manifest
                .files_for_range(&sym_spec.symbol, req_start, req_end);
            if files.is_empty() {
                debug!("No cached data for {} -- skipping", sym_spec.symbol);
                continue;
            }

            for file_path in &files {
                let remaining_capacity = MAX_FEATURES.saturating_sub(all_features.len());
                if remaining_capacity == 0 {
                    warn!(
                        "Feature buffer at capacity ({} vectors), skipping remaining files",
                        MAX_FEATURES
                    );
                    break;
                }

                // Data loading via DBN pipeline (parquet_utils removed)
                warn!(
                    "Skipping {:?}: data pipeline requires DBN loader (parquet_utils removed)",
                    file_path
                );
            }
            loaded_symbols.push(sym_spec.symbol.clone());

            if all_features.len() >= MAX_FEATURES {
                warn!(
                    "Feature buffer at capacity ({} vectors), skipping remaining symbols",
                    MAX_FEATURES
                );
                break;
            }
        }

        if all_features.is_empty() {
            return Err(MLError::ConfigError(format!(
                    "No data loaded for dataset '{}'. Check cache at {:?}",
                    spec.name, self.cache_dir
                )));
        }

        let total_bars = all_features.len();
        let (train, val, test) = split_features(all_features, &spec.split_ratio);

        Ok(DataSplit {
            name: spec.name.clone(),
            train,
            val,
            test,
            total_bars,
            symbols: loaded_symbols,
        })
    }

    /// Register a freshly-downloaded file in the cache manifest and persist it.
    pub fn register_cached_file(
        &mut self,
        symbol: &str,
        exchange: &str,
        bar_size: BarSize,
        range: CachedRange,
    ) -> Result<(), MLError> {
        self.manifest.add_range(symbol, exchange, bar_size, range);
        self.manifest.save(&self.cache_dir)
    }
}

// ---------------------------------------------------------------------------
// Debug (avoids dumping the full manifest)
// ---------------------------------------------------------------------------

impl std::fmt::Debug for DatasetManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DatasetManager")
            .field("cache_dir", &self.cache_dir)
            .field("symbols_cached", &self.manifest.symbols.len())
            .finish()
    }
}

// ---------------------------------------------------------------------------
// Splitting helper
// ---------------------------------------------------------------------------

/// Chronological split of a feature vector into train / val / test.
///
/// The split boundaries are computed from `ratio` and the total length.  Uses
/// safe slice accessors (`.get()`) to satisfy the `deny(indexing_slicing)` lint.
fn split_features(
    features: Vec<[f64; FEATURE_DIM]>,
    ratio: &SplitRatio,
) -> (
    Vec<[f64; FEATURE_DIM]>,
    Vec<[f64; FEATURE_DIM]>,
    Vec<[f64; FEATURE_DIM]>,
) {
    let n = features.len();
    let train_end = (n as f64 * ratio.train) as usize;
    let val_end = train_end + (n as f64 * ratio.val) as usize;

    let train = features.get(..train_end).unwrap_or_default().to_vec(); // cpu-side split
    let val = features
        .get(train_end..val_end)
        .unwrap_or_default()
        .to_vec(); // cpu-side split
    let test = features.get(val_end..).unwrap_or_default().to_vec(); // cpu-side split

    (train, val, test)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    /// Helper: build `n` feature vectors where `f[0] == i as f64`.
    fn make_features(n: usize) -> Vec<[f64; FEATURE_DIM]> {
        (0..n)
            .map(|i| {
                let mut f = [0.0_f64; FEATURE_DIM];
                if let Some(v) = f.get_mut(0) {
                    *v = i as f64;
                }
                f
            })
            .collect()
    }

    // -- split_features ---------------------------------------------------

    #[test]
    fn test_split_features_default_ratio() {
        let features = make_features(100);
        let ratio = SplitRatio::default();
        let (train, val, test) = split_features(features, &ratio);
        assert_eq!(train.len(), 70);
        assert_eq!(val.len(), 15);
        assert_eq!(test.len(), 15);
    }

    #[test]
    fn test_split_features_preserves_order() {
        let features = make_features(10);
        let ratio = SplitRatio {
            train: 0.5,
            val: 0.3,
            test: 0.2,
        };
        let (train, val, test) = split_features(features, &ratio);

        // Verify chronological order by inspecting f[0].
        let train_first = train
            .first()
            .and_then(|f| f.get(0).copied())
            .unwrap_or(-1.0);
        let val_first = val
            .first()
            .and_then(|f| f.get(0).copied())
            .unwrap_or(-1.0);
        let test_first = test
            .first()
            .and_then(|f| f.get(0).copied())
            .unwrap_or(-1.0);

        assert!(train_first.abs() < 1e-10);
        assert!((val_first - 5.0).abs() < 1e-10);
        assert!((test_first - 8.0).abs() < 1e-10);
    }

    #[test]
    fn test_split_features_empty() {
        let features: Vec<[f64; FEATURE_DIM]> = Vec::new();
        let ratio = SplitRatio::default();
        let (train, val, test) = split_features(features, &ratio);
        assert!(train.is_empty());
        assert!(val.is_empty());
        assert!(test.is_empty());
    }

    #[test]
    fn test_split_features_single_element() {
        let features = make_features(1);
        let ratio = SplitRatio::default();
        let (train, val, test) = split_features(features, &ratio);
        // 1 * 0.7 = 0 (truncated), so everything lands in test
        let total = train.len() + val.len() + test.len();
        assert_eq!(total, 1);
    }

    // -- DatasetManager creation ------------------------------------------

    #[test]
    fn test_dataset_manager_creation() {
        let tmp = std::env::temp_dir().join("foxhunt_manager_test");
        let _ = std::fs::remove_dir_all(&tmp);
        let manager = DatasetManager::new(tmp.clone());
        assert!(manager.is_ok());
        if let Ok(m) = manager {
            assert!(m.manifest().symbols.is_empty());
        }
        let _ = std::fs::remove_dir_all(&tmp);
    }

    // -- prepare with no data -> error ------------------------------------

    #[test]
    fn test_prepare_no_data_returns_error() {
        let tmp = std::env::temp_dir().join("foxhunt_prepare_test");
        let _ = std::fs::remove_dir_all(&tmp);
        let manager = DatasetManager::new(tmp.clone());
        assert!(manager.is_ok());

        if let Ok(m) = manager {
            let spec = DatasetSpec::from_mode(
                "empty",
                vec![super::super::SymbolSpec {
                    symbol: "AAPL".into(),
                    exchange: "XNAS".into(),
                    asset_class: super::super::AssetClass::Equity,
                }],
                super::super::DatasetMode::Dev,
            );

            let result = m.prepare(&spec);
            assert!(result.is_err());
        }
        let _ = std::fs::remove_dir_all(&tmp);
    }

    // -- Debug impl -------------------------------------------------------

    #[test]
    fn test_debug_impl() {
        let tmp = std::env::temp_dir().join("foxhunt_debug_test");
        let _ = std::fs::remove_dir_all(&tmp);
        if let Ok(m) = DatasetManager::new(tmp.clone()) {
            let debug_str = format!("{:?}", m);
            assert!(debug_str.contains("DatasetManager"));
            assert!(debug_str.contains("symbols_cached"));
        }
        let _ = std::fs::remove_dir_all(&tmp);
    }
}
