#!/usr/bin/env python3
"""
Download 90 days of ES.FUT data from Databento (Aug 3 - Nov 1, 2024).

This script downloads a continuous 90-day period of E-mini S&P 500 futures data
for unbiased DQN evaluation testing.

Usage:
    python3 scripts/python/data/download_es_90d_unseen.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"
SYMBOL = "ES.FUT"
SCHEMA = "ohlcv-1m"
DATASET = "GLBX.MDP3"

# Date range: 90 days (Aug 3 - Nov 1, 2024)
START_DATE = "2024-08-03"
END_DATE = "2024-11-01"
OUTPUT_FILE = "ES_FUT_unseen_90d.dbn"


def main():
    """Download 90 days of ES.FUT data."""
    print("=" * 80)
    print("ES.FUT 90-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"📅 Date range: {START_DATE} to {END_DATE} (91 days)")
    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)

    print()
    print("-" * 80)
    print(f"📥 Downloading 90-day range: {START_DATE} to {END_DATE}")
    print("-" * 80)

    try:
        # Parse dates
        start_dt = datetime.strptime(START_DATE, "%Y-%m-%d").replace(
            hour=0, minute=0, second=0, tzinfo=timezone.utc
        )
        end_dt = datetime.strptime(END_DATE, "%Y-%m-%d").replace(
            hour=23, minute=59, second=59, tzinfo=timezone.utc
        )

        # Build output path
        output_path = os.path.join(OUTPUT_DIR, OUTPUT_FILE)

        print(f"  Start: {start_dt.isoformat()}")
        print(f"  End:   {end_dt.isoformat()}")
        print(f"  Output: {output_path}")
        print()

        # Download data
        print("⏳ Downloading... (this may take 2-5 minutes for 90 days)")
        data = client.timeseries.get_range(
            dataset=DATASET,
            symbols=[SYMBOL],
            schema=SCHEMA,
            start=start_dt.isoformat(),
            end=end_dt.isoformat(),
        )

        # Write to file
        print("💾 Writing to file...")
        data.to_file(output_path)

        # Get file size
        file_size = os.path.getsize(output_path)
        file_size_mb = file_size / (1024 * 1024)

        print()
        print("✅ Download complete!")
        print(f"  File: {output_path}")
        print(f"  Size: {file_size:,} bytes ({file_size_mb:.2f} MB)")
        print()

        # Verify data with databento
        try:
            store = db.DBNStore.from_file(output_path)
            df = store.to_df()
            record_count = len(df)

            print("📊 Data Summary:")
            print(f"  Total bars: {record_count:,}")
            print(f"  Expected bars (90 days × ~150 bars/day): ~13,500")

            if record_count > 0:
                print(f"  Date range: {df.index[0]} to {df.index[-1]}")
                print(f"  Price range: ${df['close'].min():.2f} - ${df['close'].max():.2f}")
                print(f"  Total volume: {df['volume'].sum():,.0f}")

                # Market balance check
                bullish_bars = (df['close'] > df['open']).sum()
                bullish_pct = 100 * bullish_bars / record_count
                trend_pct = 100 * (df['close'].iloc[-1] - df['close'].iloc[0]) / df['close'].iloc[0]

                print(f"  Bullish bars: {bullish_pct:.1f}%")
                print(f"  Overall trend: {trend_pct:+.2f}%")

                # Assess balance
                if 40 <= bullish_pct <= 60:
                    print(f"  ✅ Market balance: GOOD (40-60% range)")
                elif 30 <= bullish_pct <= 70:
                    print(f"  ⚠️  Market balance: ACCEPTABLE (30-70% range)")
                else:
                    print(f"  ❌ Market balance: BIASED (outside 30-70% range)")
            else:
                print("  ⚠️  WARNING: No records in file!")

        except Exception as e:
            print(f"⚠️  Could not verify data with databento: {e}")
            print("   (File downloaded but verification failed)")

        # Estimate cost
        estimated_cost = 0.10 * 91  # ~$0.10 per day
        print()
        print(f"💰 Estimated cost: ${estimated_cost:.2f}")
        print()

        print("=" * 80)
        print("📋 NEXT STEPS")
        print("=" * 80)
        print()
        print("1. Convert DBN to Parquet:")
        print(f"   cargo run -p data --example convert_dbn_to_parquet --release -- \\")
        print(f"     --input {output_path} \\")
        print(f"     --output test_data")
        print()
        print("2. Rename output file:")
        print(f"   # Converter outputs to test_data/GLBX.MDP3/ES.FUT/...")
        print(f"   # Move and rename to test_data/ES_FUT_unseen_90d.parquet")
        print()
        print("3. Validate with DQN evaluation:")
        print("   ./test_dqn_evaluation.sh")
        print()
        print("✅ SUCCESS: 90-day ES.FUT data downloaded!")

    except Exception as e:
        print()
        print(f"❌ Download failed: {e}")
        print()
        print("Possible issues:")
        print("  • API key invalid or expired")
        print("  • Databento API rate limit exceeded")
        print("  • Network connectivity issues")
        print("  • Data not available for requested date range")
        sys.exit(1)


if __name__ == "__main__":
    main()