foxhunt/data/examples/download_mbp10_data.rs

//! Download MBP-10 (Market By Price, 10 levels) data from Databento
//!
//! This example downloads Level-2 order book data (MBP-10) for ES.FUT
//! to support TLOB neural network training.
//!
//! ## Usage
//!
//! ```bash
//! export DATABENTO_API_KEY="your_api_key_here"
//! cargo run --example download_mbp10_data --release
//! ```
//!
//! ## Data Specifications
//!
//! - **Dataset**: GLBX.MDP3 (CME Globex MDP 3.0)
//! - **Schema**: mbp-10 (Market By Price, 10 levels)
//! - **Symbols**: ES.c.0 (E-mini S&P 500 continuous front month)
//! - **Date Range**: 2024-01-02 to 2024-01-10 (7 trading days)
//! - **Encoding**: dbn (Databento Binary Encoding v2)
//! - **Compression**: zstd (Zstandard)
//! - **Expected Size**: ~50-70 GB compressed

use anyhow::{Context, Result};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::env;
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use std::time::Duration;
use tokio::time::sleep;

#[derive(Debug, Deserialize)]
struct BatchSubmitResponse {
    id: String,
}

#[derive(Debug, Deserialize)]
struct BatchStatusResponse {
    id: String,
    state: String,
    record_count: Option<u64>,
    actual_size: Option<u64>,
}

#[derive(Debug, Deserialize)]
struct BatchFileUrls {
    https: String,
    ftp: String,
}

#[derive(Debug, Deserialize)]
struct BatchFileInfo {
    filename: String,
    size: u64,
    hash: String,
    urls: BatchFileUrls,
}

const DATABENTO_API_BASE: &str = "https://hist.databento.com/v0";

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    // Get API key from environment
    let api_key =
        env::var("DATABENTO_API_KEY").context("DATABENTO_API_KEY environment variable not set")?;

    // Create output directory
    let output_dir = PathBuf::from("test_data/mbp10");
    std::fs::create_dir_all(&output_dir).context("Failed to create output directory")?;

    println!("=== Databento MBP-10 Data Download ===\n");
    println!("Dataset: GLBX.MDP3");
    println!("Schema: mbp-10 (Market By Price, 10 levels)");
    println!("Symbol: ES.c.0");
    println!("Date Range: 2024-01-02 to 2024-01-10 (7 trading days)");
    println!("Output: {}", output_dir.display());
    println!("\n{}", "=".repeat(50));

    // Step 1: Submit batch download job
    println!("\n[1/4] Submitting batch download job...");
    let job_id = submit_batch_job(&api_key).await?;
    println!("✓ Job submitted successfully: {}", job_id);

    // Step 2: Poll for job completion
    println!("\n[2/4] Waiting for job to complete...");
    poll_job_status(&api_key, &job_id).await?;
    println!("✓ Job completed successfully");

    // Step 3: Get file list and download all files
    println!("\n[3/4] Getting file list...");
    let files = list_batch_files(&api_key, &job_id).await?;
    let data_files: Vec<_> = files
        .iter()
        .filter(|f| f.filename.ends_with(".dbn.zst"))
        .collect();
    println!("✓ Found {} data files", data_files.len());

    println!("\n[4/4] Downloading data files...");
    for (i, file) in data_files.iter().enumerate() {
        println!("\n  [{}/{}] {}", i + 1, data_files.len(), file.filename);
        let output_file = output_dir.join(&file.filename);
        download_file(&api_key, &file.urls.https, &output_file).await?;
    }
    println!("✓ Download completed successfully");

    println!("\n{}", "=".repeat(50));
    println!("SUCCESS: MBP-10 data downloaded to {}", output_dir.display());
    println!("\nNext Steps:");
    println!("1. Decompress files: for f in {}/*.dbn.zst; do zstd -d $f; done", output_dir.display());
    println!("2. Validate schema: cargo run --example validate_dbn_schema");
    println!("3. Implement MBP-10 parser in data/src/providers/databento/");

    Ok(())
}

async fn submit_batch_job(api_key: &str) -> Result<String> {
    let client = Client::new();

    // Build form data (API expects application/x-www-form-urlencoded, not JSON)
    let form_data = [
        ("dataset", "GLBX.MDP3"),
        ("schema", "mbp-10"),
        ("symbols", "ES.c.0"),
        ("stype_in", "continuous"),
        ("start", "2024-01-02"),
        ("end", "2024-01-10"),
        ("encoding", "dbn"),
        ("compression", "zstd"),
    ];

    let url = format!("{}/batch.submit_job", DATABENTO_API_BASE);
    let response = client
        .post(&url)
        .basic_auth(api_key, Some(""))
        .form(&form_data)
        .send()
        .await
        .context("Failed to submit batch job")?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("API request failed: {} - {}", status, body);
    }

    let result: BatchSubmitResponse = response
        .json()
        .await
        .context("Failed to parse batch submit response")?;

    Ok(result.id)
}

async fn poll_job_status(api_key: &str, job_id: &str) -> Result<()> {
    let client = Client::new();
    let url = format!("{}/batch.list_jobs", DATABENTO_API_BASE);

    loop {
        let response = client
            .get(&url)
            .basic_auth(api_key, Some(""))
            .send()
            .await
            .context("Failed to check job status")?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            anyhow::bail!("API request failed: {} - {}", status, body);
        }

        let jobs: Vec<BatchStatusResponse> = response
            .json()
            .await
            .context("Failed to parse batch list response")?;

        // Find our job in the list
        let job = jobs
            .iter()
            .find(|j| j.id == job_id)
            .context("Job not found in batch list")?;

        match job.state.as_str() {
            "done" => {
                if let Some(size) = job.actual_size {
                    let size_mb = size as f64 / 1_048_576.0;
                    println!("  • Size: {:.2} MB", size_mb);
                }
                if let Some(records) = job.record_count {
                    println!("  • Records: {}", records);
                }
                return Ok(());
            },
            "error" => {
                anyhow::bail!("Job failed with error state");
            },
            state => {
                print!("\r  • Status: {} ... ", state);
                std::io::stdout().flush()?;
                sleep(Duration::from_secs(5)).await;
            },
        }
    }
}

async fn list_batch_files(api_key: &str, job_id: &str) -> Result<Vec<BatchFileInfo>> {
    let client = Client::new();
    let url = format!("{}/batch.list_files?job_id={}", DATABENTO_API_BASE, job_id);

    let response = client
        .get(&url)
        .basic_auth(api_key, Some(""))
        .send()
        .await
        .context("Failed to list batch files")?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("API request failed: {} - {}", status, body);
    }

    let files: Vec<BatchFileInfo> = response.json().await.context("Failed to parse file list")?;
    Ok(files)
}

async fn download_file(api_key: &str, download_url: &str, output_path: &PathBuf) -> Result<()> {
    let client = Client::new();

    let response = client
        .get(download_url)
        .basic_auth(api_key, Some(""))
        .send()
        .await
        .context("Failed to start download")?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("Download failed: {} - {}", status, body);
    }

    let total_size = response.content_length().unwrap_or(0);
    let mut file = File::create(output_path).context("Failed to create output file")?;

    let mut downloaded: u64 = 0;

    // Use chunk() method for reqwest 0.12 (replaces bytes_stream())
    let mut response = response;
    while let Some(chunk) = response.chunk().await.context("Failed to read chunk")? {
        file.write_all(&chunk)
            .context("Failed to write chunk to file")?;

        downloaded += chunk.len() as u64;

        if total_size > 0 {
            let percent = (downloaded as f64 / total_size as f64) * 100.0;
            print!(
                "\r  • Progress: {:.2}% ({:.2} MB / {:.2} MB)",
                percent,
                downloaded as f64 / 1_048_576.0,
                total_size as f64 / 1_048_576.0
            );
            std::io::stdout().flush()?;
        }
    }

    println!();
    Ok(())
}

fn validate_file(path: &PathBuf) -> Result<()> {
    let metadata = std::fs::metadata(path).context("Failed to read file metadata")?;

    let size_mb = metadata.len() as f64 / 1_048_576.0;
    println!("  • File size: {:.2} MB", size_mb);

    if size_mb < 1.0 {
        anyhow::bail!("File is suspiciously small (< 1 MB), may be corrupted");
    }

    println!("  • File appears valid");
    Ok(())
}