foxhunt/data/src/dbn_uploader.rs

//! DBN File Uploader for MinIO
//!
//! Automatically monitors test_data/real/databento/ for new .dbn files,
//! compresses them with gzip, checks for duplicates, and uploads to MinIO.
//!
//! # Features
//! - File watching with configurable poll interval
//! - Gzip compression before upload
//! - Deduplication (checks if file already exists in MinIO)
//! - Metadata tagging (symbol, schema, date range, file size)

use flate2::write::GzEncoder;
use flate2::Compression;
use std::collections::HashMap;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use tokio::fs;
use tokio::sync::RwLock;
use tokio::time::interval;
use tracing::{debug, error, info};

use crate::error::{DataError, Result as DataResult};

/// Configuration for DBN uploader
#[derive(Debug, Clone)]
pub struct DbnUploaderConfig {
    /// Directory to watch for new DBN files
    pub watch_path: PathBuf,
    /// MinIO bucket name
    pub bucket_name: String,
    /// Prefix for uploaded files (e.g., "training-data/")
    pub upload_prefix: String,
    /// Polling interval for new files
    pub poll_interval: Duration,
    /// Enable gzip compression before upload
    pub compression_enabled: bool,
    /// Enable deduplication check
    pub deduplication_enabled: bool,
}

impl Default for DbnUploaderConfig {
    fn default() -> Self {
        Self {
            watch_path: PathBuf::from("test_data/real/databento"),
            bucket_name: "ml-models".to_string(),
            upload_prefix: "training-data/".to_string(),
            poll_interval: Duration::from_secs(60),
            compression_enabled: true,
            deduplication_enabled: true,
        }
    }
}

/// Metadata extracted from DBN filename
#[derive(Debug, Clone, PartialEq)]
pub struct DbnMetadata {
    /// Trading symbol (e.g., "ES.FUT")
    pub symbol: String,
    /// Data schema (e.g., "ohlcv-1m")
    pub schema: String,
    /// Date range (e.g., "2024-01-02" or "2024-01-02_to_2024-01-31")
    pub date_range: String,
    /// Original file size in bytes
    pub file_size_bytes: u64,
}

impl DbnMetadata {
    /// Extract metadata from DBN filename
    ///
    /// Expected format: `{SYMBOL}_{SCHEMA}_{DATE_RANGE}.dbn`
    pub fn from_filename(filename: &str) -> DataResult<Self> {
        if !filename.ends_with(".dbn") {
            return Err(DataError::Validation {
                field: "filename".to_string(),
                message: format!("File must have .dbn extension: {}", filename),
            });
        }

        let name = filename.trim_end_matches(".dbn");
        let parts: Vec<&str> = name.split('_').collect();

        if parts.len() < 3 {
            return Err(DataError::Validation {
                field: "filename".to_string(),
                message: format!(
                    "Invalid DBN filename format (expected SYMBOL_SCHEMA_DATE): {}",
                    filename
                ),
            });
        }

        let symbol = parts[0].to_string();
        let schema = parts[1].to_string();
        let date_range = parts[2..].join("_");

        Ok(Self {
            symbol,
            schema,
            date_range,
            file_size_bytes: 0,
        })
    }

    /// Extract metadata from file (includes size)
    pub async fn from_file(path: &Path) -> DataResult<Self> {
        let filename = path
            .file_name()
            .ok_or_else(|| DataError::Validation {
                field: "path".to_string(),
                message: "Path has no filename".to_string(),
            })?
            .to_str()
            .ok_or_else(|| DataError::Validation {
                field: "filename".to_string(),
                message: "Filename is not valid UTF-8".to_string(),
            })?;

        let mut metadata = Self::from_filename(filename)?;
        let file_metadata = fs::metadata(path).await?;
        metadata.file_size_bytes = file_metadata.len();

        Ok(metadata)
    }
}

/// DBN file uploader
pub struct DbnUploader {
    config: DbnUploaderConfig,
    detected_files: Arc<RwLock<Vec<PathBuf>>>,
    /// Optional storage backend for MinIO/S3 uploads.
    /// When `None`, upload operations will return an error.
    storage_backend: Option<Arc<dyn storage::Storage>>,
}

impl DbnUploader {
    /// Create new uploader without a storage backend.
    ///
    /// The uploader can scan and compress files, but `upload_file` will return
    /// an error unless a storage backend is attached via `with_storage`.
    pub async fn new(config: DbnUploaderConfig) -> DataResult<Self> {
        if !config.watch_path.exists() {
            return Err(DataError::Validation {
                field: "watch_path".to_string(),
                message: format!("Watch path does not exist: {:?}", config.watch_path),
            });
        }

        info!(
            "Initializing DBN uploader: watch_path={:?}, bucket={}, prefix={}",
            config.watch_path, config.bucket_name, config.upload_prefix
        );

        Ok(Self {
            config,
            detected_files: Arc::new(RwLock::new(Vec::new())),
            storage_backend: None,
        })
    }

    /// Attach a storage backend (MinIO/S3) for actual uploads.
    ///
    /// Without a storage backend, `upload_file` and deduplication checks will
    /// return errors indicating no backend is configured.
    pub fn with_storage(mut self, backend: Arc<dyn storage::Storage>) -> Self {
        self.storage_backend = Some(backend);
        self
    }

    /// Get list of detected files (for testing)
    pub async fn get_detected_files(&self) -> Vec<PathBuf> {
        self.detected_files.read().await.clone()
    }

    /// Scan directory and return files immediately (for testing)
    ///
    /// This method is primarily for testing purposes, as `start_watching()` runs
    /// in a blocking loop. In production, use `start_watching()` which continuously
    /// monitors the directory.
    pub async fn scan_for_testing(&self) -> DataResult<Vec<PathBuf>> {
        self.scan_directory().await
    }

    /// Scan directory for DBN files
    async fn scan_directory(&self) -> DataResult<Vec<PathBuf>> {
        let mut dbn_files = Vec::new();
        let mut entries = fs::read_dir(&self.config.watch_path).await?;

        while let Some(entry) = entries.next_entry().await? {
            let path = entry.path();
            if path.extension().and_then(|s| s.to_str()) == Some("dbn") {
                debug!("Detected DBN file: {:?}", path);
                dbn_files.push(path);
            }
        }

        Ok(dbn_files)
    }

    /// Start watching for new files (blocking)
    pub async fn start_watching(&self) -> DataResult<()> {
        info!("Starting DBN file watcher...");
        let mut ticker = interval(self.config.poll_interval);

        loop {
            ticker.tick().await;

            match self.scan_directory().await {
                Ok(files) => {
                    let mut detected = self.detected_files.write().await;
                    *detected = files;
                    debug!("Scanned directory, found {} DBN files", detected.len());
                },
                Err(e) => {
                    error!("Failed to scan directory: {}", e);
                },
            }
        }
    }

    /// Check if file should be uploaded (deduplication).
    ///
    /// When deduplication is enabled and a storage backend is configured,
    /// this queries the backend to see if the upload key already exists.
    /// Returns `true` if the file should be uploaded, `false` if it already
    /// exists in the remote store.
    ///
    /// When no storage backend is configured, deduplication is skipped and
    /// the file is always considered uploadable.
    pub async fn should_upload_file(&self, path: &Path) -> DataResult<bool> {
        if !self.config.deduplication_enabled {
            return Ok(true);
        }

        let Some(backend) = &self.storage_backend else {
            // No backend configured -- cannot check remote, allow upload attempt
            debug!("No storage backend configured, skipping deduplication check");
            return Ok(true);
        };

        let key = Self::generate_upload_key(path, &self.config.upload_prefix)?;

        match backend.exists(&key).await {
            Ok(true) => {
                info!("File already exists in remote storage, skipping: {}", key);
                Ok(false)
            },
            Ok(false) => Ok(true),
            Err(e) => {
                error!("Failed to check deduplication in remote storage: {}", e);
                // On error, allow the upload attempt rather than silently skipping
                Ok(true)
            },
        }
    }

    /// Compress file with gzip
    pub async fn compress_file(path: &Path) -> DataResult<Vec<u8>> {
        let data = fs::read(path).await?;
        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(&data)?;
        Ok(encoder.finish()?)
    }

    /// Generate MinIO upload key from a file path and prefix.
    ///
    /// Returns the key in the form `{prefix}{filename}.gz`.
    ///
    /// # Errors
    ///
    /// Returns `DataError::Validation` if the path has no filename or the
    /// filename is not valid UTF-8.
    pub fn generate_upload_key(path: &Path, prefix: &str) -> DataResult<String> {
        let filename = path
            .file_name()
            .ok_or_else(|| DataError::Validation {
                field: "path".to_string(),
                message: format!("Path has no filename component: {:?}", path),
            })?
            .to_str()
            .ok_or_else(|| DataError::Validation {
                field: "filename".to_string(),
                message: format!("Filename is not valid UTF-8: {:?}", path),
            })?;
        Ok(format!("{}{}.gz", prefix, filename))
    }

    /// Generate metadata tags for MinIO
    pub fn generate_metadata_tags(metadata: &DbnMetadata) -> HashMap<String, String> {
        let mut tags = HashMap::new();
        tags.insert("symbol".to_string(), metadata.symbol.clone());
        tags.insert("schema".to_string(), metadata.schema.clone());
        tags.insert("date_range".to_string(), metadata.date_range.clone());
        tags.insert(
            "original_size".to_string(),
            metadata.file_size_bytes.to_string(),
        );
        tags
    }

    /// Upload file to MinIO/S3 with metadata.
    ///
    /// Reads the file, optionally compresses it, checks for duplicates, and
    /// stores it via the configured storage backend. The upload key is derived
    /// from the file name and the configured prefix.
    ///
    /// # Errors
    ///
    /// Returns `DataError::Storage` if no storage backend is configured, or
    /// any storage/IO error that occurs during the upload.
    pub async fn upload_file(&self, path: &Path) -> DataResult<()> {
        info!("Uploading file: {:?}", path);

        let backend = self
            .storage_backend
            .as_ref()
            .ok_or_else(|| {
                DataError::Storage(
                    "No storage backend configured. Call with_storage() to attach a \
                     MinIO/S3 backend before uploading."
                        .to_string(),
                )
            })?
            .clone();

        let metadata = DbnMetadata::from_file(path).await?;

        // Check deduplication before doing expensive compression
        if !self.should_upload_file(path).await? {
            info!("Skipping duplicate file: {:?}", path);
            return Ok(());
        }

        let data = if self.config.compression_enabled {
            Self::compress_file(path).await?
        } else {
            fs::read(path).await?
        };

        let key = Self::generate_upload_key(path, &self.config.upload_prefix)?;
        let tags = Self::generate_metadata_tags(&metadata);

        info!(
            "Uploading to storage: key={}, size={} bytes, tags={:?}",
            key,
            data.len(),
            tags
        );

        backend.store(&key, &data).await.map_err(|e| {
            DataError::Storage(format!("Failed to upload {} to storage backend: {}", key, e))
        })?;

        info!(
            "Successfully uploaded {} ({} bytes, original {} bytes)",
            key,
            data.len(),
            metadata.file_size_bytes,
        );

        Ok(())
    }
}

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

    #[tokio::test]
    async fn test_metadata_from_filename_simple() {
        let metadata =
            DbnMetadata::from_filename("ES.FUT_ohlcv-1m_2024-01-02.dbn").expect("Failed to parse");
        assert_eq!(metadata.symbol, "ES.FUT");
        assert_eq!(metadata.schema, "ohlcv-1m");
        assert_eq!(metadata.date_range, "2024-01-02");
    }

    #[tokio::test]
    async fn test_metadata_from_filename_date_range() {
        let metadata = DbnMetadata::from_filename("ZN.FUT_ohlcv-1m_2024-01-02_to_2024-01-31.dbn")
            .expect("Failed to parse");
        assert_eq!(metadata.symbol, "ZN.FUT");
        assert_eq!(metadata.schema, "ohlcv-1m");
        assert_eq!(metadata.date_range, "2024-01-02_to_2024-01-31");
    }

    #[tokio::test]
    async fn test_metadata_from_filename_invalid() {
        let result = DbnMetadata::from_filename("invalid.txt");
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_generate_upload_key() {
        let path = PathBuf::from("ES.FUT_ohlcv-1m_2024-01-02.dbn");
        let key = DbnUploader::generate_upload_key(&path, "training-data/")
            .expect("Failed to generate key");
        assert_eq!(key, "training-data/ES.FUT_ohlcv-1m_2024-01-02.dbn.gz");
    }

    #[tokio::test]
    async fn test_generate_upload_key_no_filename() {
        let path = PathBuf::from("/");
        let result = DbnUploader::generate_upload_key(&path, "prefix/");
        assert!(result.is_err(), "Should fail for path with no filename");
    }

    #[tokio::test]
    async fn test_generate_metadata_tags() {
        let metadata = DbnMetadata {
            symbol: "ES.FUT".to_string(),
            schema: "ohlcv-1m".to_string(),
            date_range: "2024-01-02".to_string(),
            file_size_bytes: 1024,
        };

        let tags = DbnUploader::generate_metadata_tags(&metadata);
        assert_eq!(
            tags.get("symbol").cloned().as_deref(),
            Some("ES.FUT")
        );
        assert_eq!(
            tags.get("schema").cloned().as_deref(),
            Some("ohlcv-1m")
        );
        assert_eq!(
            tags.get("date_range").cloned().as_deref(),
            Some("2024-01-02")
        );
        assert_eq!(
            tags.get("original_size").cloned().as_deref(),
            Some("1024")
        );
    }

    #[tokio::test]
    async fn test_upload_file_without_backend_returns_error() {
        let temp_dir = tempfile::TempDir::new().expect("Failed to create temp dir");
        let test_file = temp_dir.path().join("ES.FUT_ohlcv-1m_2024-01-02.dbn");
        fs::write(&test_file, b"test data")
            .await
            .expect("Failed to write");

        let config = DbnUploaderConfig {
            watch_path: temp_dir.path().to_path_buf(),
            compression_enabled: false,
            deduplication_enabled: false,
            ..DbnUploaderConfig::default()
        };

        let uploader = DbnUploader::new(config)
            .await
            .expect("Failed to create uploader");

        let result = uploader.upload_file(&test_file).await;
        assert!(result.is_err(), "Should fail without storage backend");
        let err_msg = format!("{}", result.unwrap_err());
        assert!(
            err_msg.contains("No storage backend configured"),
            "Error should mention missing backend, got: {}",
            err_msg
        );
    }

    #[tokio::test]
    async fn test_upload_file_with_local_storage_backend() {
        let watch_dir = tempfile::TempDir::new().expect("Failed to create watch dir");
        let storage_dir = tempfile::TempDir::new().expect("Failed to create storage dir");

        let test_file = watch_dir.path().join("ES.FUT_ohlcv-1m_2024-01-02.dbn");
        fs::write(&test_file, b"fake dbn content for upload test")
            .await
            .expect("Failed to write");

        // Use storage::local::LocalStorage as a real backend
        let local_config = storage::local::LocalStorageConfig {
            base_path: storage_dir.path().to_path_buf(),
            ..Default::default()
        };
        let local_storage = storage::local::LocalStorage::new(local_config)
            .await
            .expect("Failed to create local storage");
        let backend: Arc<dyn storage::Storage> = Arc::new(local_storage);

        let config = DbnUploaderConfig {
            watch_path: watch_dir.path().to_path_buf(),
            compression_enabled: true,
            deduplication_enabled: false,
            ..DbnUploaderConfig::default()
        };

        let uploader = DbnUploader::new(config)
            .await
            .expect("Failed to create uploader")
            .with_storage(backend.clone());

        // First upload should succeed
        uploader
            .upload_file(&test_file)
            .await
            .expect("Upload should succeed");

        // Verify the file was stored
        let key = "training-data/ES.FUT_ohlcv-1m_2024-01-02.dbn.gz";
        let exists = backend.exists(key).await.expect("Exists check failed");
        assert!(exists, "Uploaded file should exist in storage backend");
    }

    #[tokio::test]
    async fn test_deduplication_with_storage_backend() {
        let watch_dir = tempfile::TempDir::new().expect("Failed to create watch dir");
        let storage_dir = tempfile::TempDir::new().expect("Failed to create storage dir");

        let test_file = watch_dir.path().join("ES.FUT_ohlcv-1m_2024-01-02.dbn");
        fs::write(&test_file, b"dbn data")
            .await
            .expect("Failed to write");

        let local_config = storage::local::LocalStorageConfig {
            base_path: storage_dir.path().to_path_buf(),
            ..Default::default()
        };
        let local_storage = storage::local::LocalStorage::new(local_config)
            .await
            .expect("Failed to create local storage");
        let backend: Arc<dyn storage::Storage> = Arc::new(local_storage);

        let config = DbnUploaderConfig {
            watch_path: watch_dir.path().to_path_buf(),
            compression_enabled: false,
            deduplication_enabled: true,
            ..DbnUploaderConfig::default()
        };

        let uploader = DbnUploader::new(config)
            .await
            .expect("Failed to create uploader")
            .with_storage(backend.clone());

        // File does not exist yet, should_upload should be true
        let should = uploader
            .should_upload_file(&test_file)
            .await
            .expect("Check failed");
        assert!(should, "Should upload new file");

        // Upload once
        uploader
            .upload_file(&test_file)
            .await
            .expect("First upload should succeed");

        // Now the file exists, should_upload should be false
        let should = uploader
            .should_upload_file(&test_file)
            .await
            .expect("Check failed");
        assert!(!should, "Should NOT upload duplicate file");
    }
}