//! Data Pipeline for ML Training
//!
//! Automates data downloading, caching, and train/val/test splitting.
//! Provides `DatasetSpec` for declarative dataset definitions and
//! `DatasetManager` for orchestrating the full pipeline.

pub mod cache;
pub mod manager;
pub mod prepared;

pub use cache::{CacheManifest, CachedRange, SymbolCache};
pub use manager::{DataSplit, DatasetManager, FEATURE_DIM};
pub use prepared::{Batch, BatchIterator, PreparedDataset};

use chrono::NaiveDate;
use serde::{Deserialize, Serialize};

// Re-export AssetClass from asset_selection (canonical location)
pub use crate::asset_selection::AssetClass;

/// Bar size / time resolution
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum BarSize {
    OneMinute,
    FiveMinute,
    FifteenMinute,
    OneHour,
    Daily,
}

impl BarSize {
    /// Estimated bars per trading day (6.5 hours for equities)
    pub fn bars_per_day(&self) -> usize {
        match self {
            BarSize::OneMinute => 390,
            BarSize::FiveMinute => 78,
            BarSize::FifteenMinute => 26,
            BarSize::OneHour => 7,
            BarSize::Daily => 1,
        }
    }
}

/// Date range specification
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum DateRange {
    /// Absolute date range
    Absolute { start: NaiveDate, end: NaiveDate },
    /// Rolling window from today
    Rolling { days: u32 },
}

impl DateRange {
    /// Resolve to absolute start/end dates
    pub fn resolve(&self) -> (NaiveDate, NaiveDate) {
        match self {
            DateRange::Absolute { start, end } => (*start, *end),
            DateRange::Rolling { days } => {
                let end = chrono::Utc::now().date_naive();
                let start = end - chrono::Duration::days(i64::from(*days));
                (start, end)
            }
        }
    }

    /// Number of calendar days in range
    pub fn calendar_days(&self) -> u32 {
        let (start, end) = self.resolve();
        (end - start).num_days().unsigned_abs() as u32
    }
}

/// Train/validation/test split ratios
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SplitRatio {
    pub train: f64,
    pub val: f64,
    pub test: f64,
}

impl Default for SplitRatio {
    fn default() -> Self {
        Self {
            train: 0.70,
            val: 0.15,
            test: 0.15,
        }
    }
}

impl SplitRatio {
    /// Validate that ratios sum to 1.0
    pub fn is_valid(&self) -> bool {
        let sum = self.train + self.val + self.test;
        (sum - 1.0).abs() < 1e-6 && self.train > 0.0 && self.val >= 0.0 && self.test >= 0.0
    }
}

/// Preset dataset modes
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum DatasetMode {
    /// Last 30 days -- fast iteration (~40K bars/symbol at 1m)
    Dev,
    /// Last 6 months -- regime diversity (~120K bars/symbol at 1m)
    Backtest,
    /// 2+ years -- production training (~480K bars/symbol at 1m)
    Full,
}

impl DatasetMode {
    /// Default rolling window days for this mode
    pub fn default_days(&self) -> u32 {
        match self {
            DatasetMode::Dev => 30,
            DatasetMode::Backtest => 180,
            DatasetMode::Full => 730,
        }
    }

    /// Create a DateRange for this mode
    pub fn to_date_range(&self) -> DateRange {
        DateRange::Rolling {
            days: self.default_days(),
        }
    }
}

/// Specification for a single symbol in the dataset
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SymbolSpec {
    /// Trading symbol (e.g., "AAPL", "SPY")
    pub symbol: String,
    /// Exchange identifier (e.g., "XNAS", "XNYS")
    pub exchange: String,
    /// Asset classification
    pub asset_class: AssetClass,
}

/// Complete dataset specification -- declarative definition of training data needs
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatasetSpec {
    /// Human-readable name (e.g., "us-equities-dev")
    pub name: String,
    /// Symbols to include
    pub symbols: Vec<SymbolSpec>,
    /// Date range for data
    pub date_range: DateRange,
    /// Bar time resolution
    pub bar_size: BarSize,
    /// Train/val/test split ratios
    pub split_ratio: SplitRatio,
    /// Warmup bars for technical indicators (default: 50)
    pub warmup_bars: usize,
}

impl DatasetSpec {
    /// Create a spec from a preset mode
    pub fn from_mode(name: &str, symbols: Vec<SymbolSpec>, mode: DatasetMode) -> Self {
        Self {
            name: name.to_string(),
            symbols,
            date_range: mode.to_date_range(),
            bar_size: BarSize::OneMinute,
            split_ratio: SplitRatio::default(),
            warmup_bars: 50,
        }
    }

    /// Create a spec from an asset universe and dataset mode.
    ///
    /// Only includes enabled assets from the universe.
    #[must_use]
    pub fn from_universe(
        name: &str,
        universe: &crate::asset_selection::AssetUniverse,
        mode: DatasetMode,
    ) -> Self {
        let symbols = universe
            .eligible()
            .iter()
            .map(|a| SymbolSpec {
                symbol: a.symbol.clone(),
                exchange: a.exchange.clone(),
                asset_class: a.asset_class.clone(),
            })
            .collect();

        Self {
            name: format!("{name}-{mode:?}"),
            symbols,
            date_range: mode.to_date_range(),
            bar_size: BarSize::OneMinute,
            split_ratio: SplitRatio::default(),
            warmup_bars: 50,
        }
    }

    /// Estimate total bars across all symbols
    pub fn estimated_total_bars(&self) -> usize {
        let trading_days = (self.date_range.calendar_days() as f64 * 252.0 / 365.0) as usize;
        let bars_per_symbol = trading_days * self.bar_size.bars_per_day();
        bars_per_symbol * self.symbols.len()
    }
}

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

    fn sample_symbols() -> Vec<SymbolSpec> {
        vec![
            SymbolSpec {
                symbol: "AAPL".into(),
                exchange: "XNAS".into(),
                asset_class: AssetClass::Equity,
            },
            SymbolSpec {
                symbol: "SPY".into(),
                exchange: "XNYS".into(),
                asset_class: AssetClass::ETF,
            },
        ]
    }

    #[test]
    fn test_dataset_mode_days() {
        assert_eq!(DatasetMode::Dev.default_days(), 30);
        assert_eq!(DatasetMode::Backtest.default_days(), 180);
        assert_eq!(DatasetMode::Full.default_days(), 730);
    }

    #[test]
    fn test_bar_size_bars_per_day() {
        assert_eq!(BarSize::OneMinute.bars_per_day(), 390);
        assert_eq!(BarSize::Daily.bars_per_day(), 1);
    }

    #[test]
    fn test_split_ratio_default_valid() {
        let ratio = SplitRatio::default();
        assert!(ratio.is_valid());
    }

    #[test]
    fn test_split_ratio_invalid() {
        let ratio = SplitRatio {
            train: 0.5,
            val: 0.5,
            test: 0.5,
        };
        assert!(!ratio.is_valid());
    }

    #[test]
    fn test_date_range_rolling_days() {
        let range = DateRange::Rolling { days: 30 };
        assert_eq!(range.calendar_days(), 30);
    }

    #[test]
    fn test_dataset_spec_from_mode() {
        let spec = DatasetSpec::from_mode("test", sample_symbols(), DatasetMode::Dev);
        assert_eq!(spec.name, "test");
        assert_eq!(spec.symbols.len(), 2);
        assert_eq!(spec.warmup_bars, 50);
        assert!(spec.split_ratio.is_valid());
    }

    #[test]
    fn test_estimated_total_bars() {
        let spec = DatasetSpec::from_mode("test", sample_symbols(), DatasetMode::Dev);
        let bars = spec.estimated_total_bars();
        // 30 days * ~252/365 trading ratio * 390 bars/day * 2 symbols
        assert!(bars > 10_000);
        assert!(bars < 100_000);
    }

    #[test]
    fn test_absolute_date_range() {
        let range = DateRange::Absolute {
            start: NaiveDate::from_ymd_opt(2024, 1, 1).unwrap_or_default(),
            end: NaiveDate::from_ymd_opt(2024, 6, 30).unwrap_or_default(),
        };
        assert_eq!(range.calendar_days(), 181);
    }

    #[test]
    fn test_asset_class_serde() {
        let json = serde_json::to_string(&AssetClass::Equity).unwrap_or_default();
        assert!(json.contains("Equity"));
    }

    #[test]
    fn test_asset_class_is_same_type() {
        let dp_future = AssetClass::Future;
        let as_future = crate::asset_selection::AssetClass::Future;
        assert_eq!(dp_future, as_future);
    }

    #[test]
    fn test_dataset_spec_from_universe() {
        let universe = crate::asset_selection::AssetUniverse::futures_baseline();
        let spec = DatasetSpec::from_universe("futures-baseline", &universe, DatasetMode::Dev);

        assert_eq!(spec.symbols.len(), 4);
        assert!(spec.name.contains("Dev"));
        assert_eq!(spec.bar_size, BarSize::OneMinute);
        assert!(spec.split_ratio.is_valid());

        // Verify symbols match universe
        let symbol_names: Vec<&str> = spec.symbols.iter().map(|s| s.symbol.as_str()).collect();
        assert!(symbol_names.contains(&"ES.FUT"));
        assert!(symbol_names.contains(&"NQ.FUT"));
        assert!(symbol_names.contains(&"ZN.FUT"));
        assert!(symbol_names.contains(&"6E.FUT"));

        // All should be Future class
        for sym in &spec.symbols {
            assert_eq!(sym.asset_class, AssetClass::Future);
        }
    }

    #[test]
    fn test_dataset_spec_from_universe_modes() {
        let universe = crate::asset_selection::AssetUniverse::futures_baseline();

        let dev = DatasetSpec::from_universe("futures-baseline", &universe, DatasetMode::Dev);
        let backtest = DatasetSpec::from_universe("futures-baseline", &universe, DatasetMode::Backtest);
        let full = DatasetSpec::from_universe("futures-baseline", &universe, DatasetMode::Full);

        // Dev < Backtest < Full in estimated bars
        assert!(dev.estimated_total_bars() < backtest.estimated_total_bars());
        assert!(backtest.estimated_total_bars() < full.estimated_total_bars());
    }

    #[test]
    fn test_dataset_spec_from_universe_disabled_filtered() {
        let mut universe = crate::asset_selection::AssetUniverse::futures_baseline();
        // Disable one asset
        if let Some(first) = universe.assets.get_mut(0) {
            first.enabled = false;
        }
        let spec = DatasetSpec::from_universe("futures-baseline", &universe, DatasetMode::Dev);
        // Only 3 symbols because one was disabled
        assert_eq!(spec.symbols.len(), 3);
    }
}
