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

This script downloads multiple days of E-mini S&P 500 futures data
using specific contract codes (ESH4, ESM4, etc.) to ensure data availability.

Usage:
    python3 download_es_databento_v2.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"
SCHEMA = "ohlcv-1m"
DATASET = "GLBX.MDP3"

# Target dates with specific contracts
# ESH4 = March 2024 contract (expires mid-March)
# ESM4 = June 2024 contract (expires mid-June)
DOWNLOAD_CONFIGS = [
    {
        "date": "2024-01-03",
        "symbol": "ESH4",
        "regime": "Trending",
        "description": "Strong uptrend continuation from Jan 2"
    },
    {
        "date": "2024-01-04",
        "symbol": "ESH4",
        "regime": "Ranging",
        "description": "Consolidation, sideways movement"
    },
    {
        "date": "2024-01-05",
        "symbol": "ESH4",
        "regime": "Volatile",
        "description": "High volatility, whipsaws"
    },
]


def main():
    """Download ES futures data for multiple days."""
    print("=" * 80)
    print("ES Futures Multi-Day Databento Download (V2)")
    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"📊 Schema: {SCHEMA}")
    print(f"📦 Dataset: {DATASET}")
    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 config in DOWNLOAD_CONFIGS:
        date_str = config["date"]
        symbol = config["symbol"]
        regime = config["regime"]
        description = config["description"]

        print()
        print("-" * 80)
        print(f"📥 Downloading: {date_str} ({regime})")
        print(f"   Symbol: {symbol}")
        print(f"   {description}")
        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

            # Read back to verify data count
            try:
                store = db.DBNStore.from_file(output_file)
                df = store.to_df()
                record_count = len(df)

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

                # Show sample data
                if record_count > 0:
                    print(f"  Price range: ${df['close'].min():.2f} - ${df['close'].max():.2f}")
                    print(f"  Volume: {df['volume'].sum():,.0f}")
                else:
                    print("  ⚠️  WARNING: No records in file!")

            except Exception as e:
                print(f"✅ Download complete!")
                print(f"  File size: {file_size:,} bytes ({file_size_kb:.2f} KB)")
                print(f"  ⚠️  Could not verify record count: {e}")

            # 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": date_str,
                "symbol": symbol,
                "regime": regime,
                "path": output_file,
                "size_kb": file_size_kb
            })

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

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

    if successful_downloads:
        print("✅ Successfully downloaded files:")
        for download in successful_downloads:
            print(f"  • {download['date']} ({download['regime']}): {download['symbol']} - {download['size_kb']:.2f} KB")
        print()

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

    # Regime classification summary
    print("📋 REGIME CLASSIFICATION:")
    for config in DOWNLOAD_CONFIGS:
        status = "✅" if any(d["date"] == config["date"] for d in successful_downloads) else "❌"
        print(f"  {status} {config['date']} - {config['regime']}: {config['description']}")
    print()

    print("📋 NEXT STEPS:")
    print("1. Validate each file:")
    print("   cd /home/jgrusewski/Work/foxhunt")
    print("   cargo run -p backtesting_service --example validate_dbn_data")
    print("2. Use data for regime detection testing in adaptive strategy")
    print("3. Analyze market characteristics:")
    print("   • Price movements, volatility patterns")
    print("   • Volume distribution")
    print("   • Regime transition detection")
    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()