foxhunt/ml/src/data_pipeline/cache.rs

//! Cache manifest for tracking downloaded training data.
//!
//! The [`CacheManifest`] records which date ranges have been downloaded for each
//! symbol so that [`super::manager::DatasetManager`] can identify gaps and only
//! fetch missing data.

use std::collections::HashMap;
use std::path::{Path, PathBuf};

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

use super::BarSize;
use crate::MLError;

// ---------------------------------------------------------------------------
// Core types
// ---------------------------------------------------------------------------

/// A single contiguous range of cached bar data for one symbol.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CachedRange {
    /// First trading date (inclusive) in this chunk.
    pub start: NaiveDate,
    /// Last trading date (inclusive) in this chunk.
    pub end: NaiveDate,
    /// Number of bars stored in the file.
    pub bar_count: usize,
    /// Path to the parquet/csv file relative to the cache directory.
    pub file_path: PathBuf,
    /// When this range was downloaded.
    pub cached_at: DateTime<Utc>,
}

/// Per-symbol cache metadata.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SymbolCache {
    /// Exchange identifier (e.g. "XNAS").
    pub exchange: String,
    /// Bar resolution that was cached.
    pub bar_size: BarSize,
    /// Cached date ranges, kept sorted by `start`.
    pub ranges: Vec<CachedRange>,
    /// Running total of bars across all ranges.
    pub total_bars: usize,
}

/// Top-level manifest persisted as `manifest.json` inside the cache directory.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheManifest {
    /// Last time the manifest was written to disk.
    pub last_updated: DateTime<Utc>,
    /// Per-symbol cache state keyed by symbol string (e.g. "AAPL").
    pub symbols: HashMap<String, SymbolCache>,
}

impl Default for CacheManifest {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// Implementation
// ---------------------------------------------------------------------------

impl CacheManifest {
    /// Create a brand-new empty manifest.
    #[must_use]
    pub fn new() -> Self {
        Self {
            last_updated: Utc::now(),
            symbols: HashMap::new(),
        }
    }

    /// Load a manifest from `<cache_dir>/manifest.json`.
    ///
    /// Returns an empty manifest if the file does not exist.
    ///
    /// # Errors
    /// Returns `MLError::ConfigError` on I/O or parse failures (other than
    /// file-not-found).
    pub fn load(cache_dir: &Path) -> Result<Self, MLError> {
        let path = cache_dir.join("manifest.json");
        match std::fs::read_to_string(&path) {
            Ok(contents) => serde_json::from_str(&contents).map_err(|e| MLError::ConfigError {
                reason: format!("Failed to parse manifest at {}: {e}", path.display()),
            }),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Self::new()),
            Err(e) => Err(MLError::ConfigError {
                reason: format!("Failed to read manifest at {}: {e}", path.display()),
            }),
        }
    }

    /// Persist the manifest to `<cache_dir>/manifest.json`.
    ///
    /// Creates the directory if it does not exist. Updates `last_updated`.
    ///
    /// # Errors
    /// Returns `MLError::ConfigError` on I/O or serialization failures.
    pub fn save(&mut self, cache_dir: &Path) -> Result<(), MLError> {
        std::fs::create_dir_all(cache_dir).map_err(|e| MLError::ConfigError {
            reason: format!("Failed to create cache dir {}: {e}", cache_dir.display()),
        })?;

        self.last_updated = Utc::now();
        let json = serde_json::to_string_pretty(self).map_err(|e| MLError::ConfigError {
            reason: format!("Failed to serialize manifest: {e}"),
        })?;

        let path = cache_dir.join("manifest.json");
        std::fs::write(&path, json).map_err(|e| MLError::ConfigError {
            reason: format!("Failed to write manifest to {}: {e}", path.display()),
        })?;

        Ok(())
    }

    /// Register a newly-downloaded range for a symbol.
    ///
    /// Keeps `ranges` sorted by start date and recalculates `total_bars`.
    pub fn add_range(
        &mut self,
        symbol: &str,
        exchange: &str,
        bar_size: BarSize,
        range: CachedRange,
    ) {
        let entry = self
            .symbols
            .entry(symbol.to_string())
            .or_insert_with(|| SymbolCache {
                exchange: exchange.to_string(),
                bar_size,
                ranges: Vec::new(),
                total_bars: 0,
            });

        entry.ranges.push(range);
        entry.ranges.sort_by_key(|r| r.start);
        entry.total_bars = entry.ranges.iter().map(|r| r.bar_count).sum();
    }

    /// Identify date gaps for `symbol` within the requested `[start, end]`.
    ///
    /// Returns a list of `(gap_start, gap_end)` pairs representing ranges that
    /// are **not** present in the cache and therefore need downloading.
    ///
    /// If the symbol has no cached data at all, the entire requested range is
    /// returned as a single gap.
    #[must_use]
    pub fn find_gaps(
        &self,
        symbol: &str,
        start: NaiveDate,
        end: NaiveDate,
    ) -> Vec<(NaiveDate, NaiveDate)> {
        let ranges = match self.symbols.get(symbol) {
            Some(sc) => &sc.ranges,
            None => return vec![(start, end)],
        };

        if ranges.is_empty() {
            return vec![(start, end)];
        }

        let mut gaps = Vec::new();
        let mut cursor = start;

        for r in ranges {
            // Skip ranges entirely before our window.
            if r.end < start {
                continue;
            }
            // Stop once ranges are past our window.
            if r.start > end {
                break;
            }

            // If there is a gap between the cursor and this range's start,
            // record it (clamped to the requested window).
            let range_start = r.start.max(start);
            if cursor < range_start {
                gaps.push((cursor, range_start));
            }

            // Advance cursor past this cached range (but not beyond `end`).
            if r.end >= cursor {
                cursor = r.end;
            }
        }

        // Trailing gap after all cached ranges.
        if cursor < end {
            gaps.push((cursor, end));
        }

        gaps
    }

    /// Return the file paths whose cached ranges overlap `[start, end]` for a
    /// given symbol.
    #[must_use]
    pub fn files_for_range(
        &self,
        symbol: &str,
        start: NaiveDate,
        end: NaiveDate,
    ) -> Vec<PathBuf> {
        let ranges = match self.symbols.get(symbol) {
            Some(sc) => &sc.ranges,
            None => return Vec::new(),
        };

        ranges
            .iter()
            .filter(|r| r.start <= end && r.end >= start)
            .map(|r| r.file_path.clone())
            .collect()
    }
}

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

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

    /// Helper: create a `NaiveDate` from y-m-d, falling back to epoch on
    /// invalid input (avoids `unwrap()`).
    fn date(y: i32, m: u32, d: u32) -> NaiveDate {
        NaiveDate::from_ymd_opt(y, m, d).unwrap_or_default()
    }

    /// Helper: build a minimal `CachedRange`.
    fn cached(start: NaiveDate, end: NaiveDate, bars: usize) -> CachedRange {
        CachedRange {
            start,
            end,
            bar_count: bars,
            file_path: PathBuf::from(format!(
                "{}_{}.parquet",
                start.format("%Y%m%d"),
                end.format("%Y%m%d")
            )),
            cached_at: Utc::now(),
        }
    }

    // -----------------------------------------------------------------------
    // 1. Empty manifest -> entire range is a gap
    // -----------------------------------------------------------------------
    #[test]
    fn empty_manifest_full_gap() {
        let manifest = CacheManifest::new();
        let gaps = manifest.find_gaps("AAPL", date(2024, 1, 1), date(2024, 6, 30));

        assert_eq!(gaps.len(), 1);
        let (gs, ge) = gaps.first().copied().unwrap_or_default();
        assert_eq!(gs, date(2024, 1, 1));
        assert_eq!(ge, date(2024, 6, 30));
    }

    // -----------------------------------------------------------------------
    // 2. Partial cache -> one leading gap
    // -----------------------------------------------------------------------
    #[test]
    fn partial_cache_finds_gap() {
        let mut manifest = CacheManifest::new();
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 4, 1), date(2024, 6, 30), 65),
        );

        let gaps = manifest.find_gaps("AAPL", date(2024, 1, 1), date(2024, 6, 30));

        assert_eq!(gaps.len(), 1);
        let (gs, ge) = gaps.first().copied().unwrap_or_default();
        assert_eq!(gs, date(2024, 1, 1));
        assert_eq!(ge, date(2024, 4, 1));
    }

    // -----------------------------------------------------------------------
    // 3. Full cache -> no gaps
    // -----------------------------------------------------------------------
    #[test]
    fn full_cache_no_gaps() {
        let mut manifest = CacheManifest::new();
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 1, 1), date(2024, 6, 30), 130),
        );

        let gaps = manifest.find_gaps("AAPL", date(2024, 1, 1), date(2024, 6, 30));
        assert!(gaps.is_empty());
    }

    // -----------------------------------------------------------------------
    // 4. Middle gap between two cached ranges
    // -----------------------------------------------------------------------
    #[test]
    fn middle_gap() {
        let mut manifest = CacheManifest::new();
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 1, 1), date(2024, 2, 28), 40),
        );
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 5, 1), date(2024, 6, 30), 44),
        );

        let gaps = manifest.find_gaps("AAPL", date(2024, 1, 1), date(2024, 6, 30));

        assert_eq!(gaps.len(), 1);
        let (gs, ge) = gaps.first().copied().unwrap_or_default();
        assert_eq!(gs, date(2024, 2, 28));
        assert_eq!(ge, date(2024, 5, 1));
    }

    // -----------------------------------------------------------------------
    // 5. files_for_range returns overlapping files
    // -----------------------------------------------------------------------
    #[test]
    fn files_for_range_returns_overlapping() {
        let mut manifest = CacheManifest::new();
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 1, 1), date(2024, 3, 31), 65),
        );
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 4, 1), date(2024, 6, 30), 65),
        );
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 7, 1), date(2024, 9, 30), 66),
        );

        // Query Feb-May should return the first two files.
        let files = manifest.files_for_range("AAPL", date(2024, 2, 1), date(2024, 5, 31));
        assert_eq!(files.len(), 2);

        // Query for a symbol with no data returns empty.
        let files = manifest.files_for_range("MSFT", date(2024, 1, 1), date(2024, 6, 30));
        assert!(files.is_empty());
    }

    // -----------------------------------------------------------------------
    // 6. save and load round-trip
    // -----------------------------------------------------------------------
    #[test]
    fn save_and_load_manifest() {
        let dir = std::env::temp_dir().join(format!("foxhunt_cache_test_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);

        let mut manifest = CacheManifest::new();
        manifest.add_range(
            "SPY",
            "XNYS",
            BarSize::FiveMinute,
            cached(date(2024, 3, 1), date(2024, 3, 31), 2340),
        );

        // Save should succeed.
        let save_result = manifest.save(&dir);
        assert!(save_result.is_ok());

        // Load should round-trip.
        let loaded = CacheManifest::load(&dir);
        assert!(loaded.is_ok());
        let loaded = loaded.unwrap_or_default();

        assert!(loaded.symbols.contains_key("SPY"));
        let spy = loaded.symbols.get("SPY");
        assert!(spy.is_some());
        if let Some(spy) = spy {
            assert_eq!(spy.ranges.len(), 1);
            assert_eq!(spy.total_bars, 2340);
            assert!(matches!(spy.bar_size, BarSize::FiveMinute));
        }

        // Cleanup.
        let _ = std::fs::remove_dir_all(&dir);
    }

    // -----------------------------------------------------------------------
    // 7. Build manifest from on-disk DBN files
    // -----------------------------------------------------------------------
    #[test]
    fn test_build_manifest_from_disk() {
        use std::path::Path;

        let cache_dir = Path::new("../data/cache");
        if !cache_dir.exists() {
            // Skip if data hasn't been moved yet
            return;
        }

        let mut manifest = CacheManifest::new();

        // Scan each symbol directory
        for symbol in &["ES.FUT", "NQ.FUT", "ZN.FUT", "6E.FUT"] {
            let symbol_dir = cache_dir.join(symbol);
            if !symbol_dir.exists() {
                continue;
            }

            let mut files: Vec<_> = std::fs::read_dir(&symbol_dir)
                .into_iter()
                .flatten()
                .filter_map(|e| e.ok())
                .filter(|e| {
                    let path = e.path();
                    let name = path.file_name().unwrap_or_default().to_string_lossy();
                    name.ends_with(".dbn") || name.ends_with(".dbn.zst")
                })
                .collect();
            files.sort_by_key(|e| e.file_name());

            for entry in &files {
                let fname = entry.file_name();
                let fname_str = fname.to_string_lossy();
                // Parse date from filename: ohlcv-1m_2024-01-02.dbn
                if let Some(date_str) = fname_str
                    .strip_prefix("ohlcv-1m_")
                    .and_then(|s| s.strip_suffix(".dbn.zst").or_else(|| s.strip_suffix(".dbn")))
                {
                    if let Ok(parsed_date) =
                        NaiveDate::parse_from_str(date_str, "%Y-%m-%d")
                    {
                        manifest.add_range(
                            symbol,
                            "GLBX.MDP3",
                            BarSize::OneMinute,
                            CachedRange {
                                start: parsed_date,
                                end: parsed_date,
                                bar_count: 390, // Approximate for 1m bars
                                file_path: entry
                                    .path()
                                    .strip_prefix(cache_dir)
                                    .unwrap_or(&entry.path())
                                    .to_path_buf(),
                                cached_at: Utc::now(),
                            },
                        );
                    }
                }
            }
        }

        // Skip if no data was found (symbols may be in a subdirectory like futures-baseline/)
        if manifest.symbols.is_empty() {
            return;
        }

        // Save the manifest
        let result = manifest.save(cache_dir);
        assert!(result.is_ok(), "Failed to save manifest: {:?}", result.err());

        // Verify it round-trips
        let loaded = CacheManifest::load(cache_dir);
        assert!(loaded.is_ok());
        let loaded = loaded.unwrap_or_default();
        assert_eq!(loaded.symbols.len(), manifest.symbols.len());
    }

    // -----------------------------------------------------------------------
    // 8. add_range keeps ranges sorted by start date
    // -----------------------------------------------------------------------
    #[test]
    fn add_range_sorts_by_start() {
        let mut manifest = CacheManifest::new();

        // Insert in reverse chronological order.
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 7, 1), date(2024, 9, 30), 66),
        );
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 1, 1), date(2024, 3, 31), 65),
        );
        manifest.add_range(
            "AAPL",
            "XNAS",
            BarSize::Daily,
            cached(date(2024, 4, 1), date(2024, 6, 30), 65),
        );

        let entry = manifest.symbols.get("AAPL");
        assert!(entry.is_some());
        if let Some(entry) = entry {
            assert_eq!(entry.ranges.len(), 3);
            assert_eq!(entry.total_bars, 196);

            // Verify sorted order.
            let first = entry.ranges.first();
            let last = entry.ranges.last();
            assert!(first.is_some());
            assert!(last.is_some());
            if let (Some(f), Some(l)) = (first, last) {
                assert_eq!(f.start, date(2024, 1, 1));
                assert_eq!(l.start, date(2024, 7, 1));
            }
        }
    }
}