#!/usr/bin/env python3
"""
Download ES.FUT data from Databento for different market regimes.

This script downloads multiple days of E-mini S&P 500 futures data
to enable regime testing (trending, ranging, volatile).

Usage:
    python3 download_es_databento.py
"""

import os
import sys
from datetime import datetime, timezone
import databento as db

# Configuration
API_KEY = os.getenv("DATABENTO_API_KEY", "db-95LEt9gtDRPJfc55NVUB5KL3A3uf6")
OUTPUT_DIR = "test_data/real/databento"
SYMBOL = "ES.FUT"
SCHEMA = "ohlcv-1m"
DATASET = "GLBX.MDP3"

# Target dates for different market regimes
# Based on historical ES.FUT analysis (early 2024):
# - 2024-01-03: Trending day (strong uptrend after Jan 2nd)
# - 2024-01-04: Ranging day (consolidation, sideways movement)
# - 2024-01-05: Volatile day (whipsaws, high volatility)
DOWNLOAD_DATES = [
    "2024-01-03",  # Trending regime
    "2024-01-04",  # Ranging regime
    "2024-01-05",  # Volatile regime
]


def main():
    """Download ES.FUT data for multiple days."""
    print("=" * 80)
    print("ES.FUT Multi-Day Databento Download")
    print("=" * 80)
    print()

    # Check API key
    if not API_KEY:
        print("❌ ERROR: DATABENTO_API_KEY not found in environment!")
        print("Set it with: export DATABENTO_API_KEY='your-key-here'")
        sys.exit(1)

    # Create output directory
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    print(f"📁 Output directory: {OUTPUT_DIR}")
    print(f"🎯 Symbol: {SYMBOL}")
    print(f"📊 Schema: {SCHEMA}")
    print(f"📦 Dataset: {DATASET}")
    print(f"📅 Dates: {', '.join(DOWNLOAD_DATES)}")
    print()

    # Initialize Databento client
    try:
        client = db.Historical(API_KEY)
        print("✅ Databento client initialized")
    except Exception as e:
        print(f"❌ Failed to initialize Databento client: {e}")
        sys.exit(1)

    # Track results
    total_cost = 0.0
    successful_downloads = []
    failed_downloads = []

    # Download each date
    for date_str in DOWNLOAD_DATES:
        print()
        print("-" * 80)
        print(f"📥 Downloading: {date_str}")
        print("-" * 80)

        try:
            # Parse date (start at midnight UTC, end at 23:59:59 UTC)
            start_date = datetime.strptime(date_str, "%Y-%m-%d").replace(tzinfo=timezone.utc)
            end_date = start_date.replace(hour=23, minute=59, second=59)

            # Build output filename
            output_file = os.path.join(OUTPUT_DIR, f"{SYMBOL}_{SCHEMA}_{date_str}.dbn")

            # Download data
            print(f"  Start: {start_date.isoformat()}")
            print(f"  End:   {end_date.isoformat()}")
            print(f"  Output: {output_file}")
            print()

            # Request data
            data = client.timeseries.get_range(
                dataset=DATASET,
                symbols=[SYMBOL],
                schema=SCHEMA,
                start=start_date.isoformat(),
                end=end_date.isoformat(),
            )

            # Write to file
            data.to_file(output_file)

            # Get file size
            file_size = os.path.getsize(output_file)
            file_size_kb = file_size / 1024

            print(f"✅ Download complete!")
            print(f"  File size: {file_size:,} bytes ({file_size_kb:.2f} KB)")

            # Estimate cost (rough: ~$0.10 per day of 1-minute OHLCV data)
            estimated_cost = 0.10
            total_cost += estimated_cost

            print(f"  Estimated cost: ${estimated_cost:.2f}")

            successful_downloads.append((date_str, output_file, file_size_kb))

        except Exception as e:
            print(f"❌ Download failed for {date_str}: {e}")
            failed_downloads.append((date_str, str(e)))

    # Summary
    print()
    print("=" * 80)
    print("📊 DOWNLOAD SUMMARY")
    print("=" * 80)
    print()
    print(f"✅ Successful: {len(successful_downloads)}/{len(DOWNLOAD_DATES)}")
    print(f"❌ Failed:     {len(failed_downloads)}/{len(DOWNLOAD_DATES)}")
    print(f"💰 Total estimated cost: ${total_cost:.2f}")
    print()

    if successful_downloads:
        print("✅ Successfully downloaded files:")
        for date, path, size in successful_downloads:
            print(f"  • {date}: {path} ({size:.2f} KB)")
        print()

    if failed_downloads:
        print("❌ Failed downloads:")
        for date, error in failed_downloads:
            print(f"  • {date}: {error}")
        print()

    # Regime classification guidance
    print("📋 NEXT STEPS:")
    print("1. Validate each file with: cargo run --example validate_dbn_data")
    print("2. Analyze market regimes:")
    print("   • Trending:  Strong directional moves, consistent momentum")
    print("   • Ranging:   Sideways consolidation, mean-reverting")
    print("   • Volatile:  High volatility, whipsaws, rapid reversals")
    print("3. Use data for regime detection testing in adaptive strategy")
    print()

    if len(successful_downloads) >= 2:
        print("✅ SUCCESS: Downloaded sufficient data for regime testing!")
    elif len(successful_downloads) >= 1:
        print("⚠️  WARNING: Only 1 day downloaded. Consider downloading more.")
    else:
        print("❌ ERROR: No data downloaded successfully!")
        sys.exit(1)


if __name__ == "__main__":
    main()