#!/usr/bin/env python3
"""
Download 90 days of ES futures data using multiple contracts (Aug 3 - Nov 1, 2024).

This script downloads data from two ES futures contracts:
- ESU4 (September 2024): Aug 3 - Sep 19
- ESZ4 (December 2024): Sep 20 - Nov 1

Then merges them into a single continuous dataset.

Usage:
    .venv/bin/python3 scripts/python/data/download_es_90d_multi_contract.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"
SCHEMA = "ohlcv-1m"
DATASET = "GLBX.MDP3"

# Contract periods
CONTRACTS = [
    {
        "symbol": "ESU4",
        "start": "2024-08-03",
        "end": "2024-09-19",
        "description": "September 2024 contract (Aug 3 - Sep 19)"
    },
    {
        "symbol": "ESZ4",
        "start": "2024-09-20",
        "end": "2024-11-01",
        "description": "December 2024 contract (Sep 20 - Nov 1)"
    },
]


def download_contract(client, contract):
    """Download a single contract's data."""
    symbol = contract["symbol"]
    start_str = contract["start"]
    end_str = contract["end"]

    print()
    print("-" * 80)
    print(f"📥 Downloading: {symbol}")
    print(f"   {contract['description']}")
    print("-" * 80)

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

        # Build output filename
        output_file = os.path.join(OUTPUT_DIR, f"{symbol}_90d_part.dbn")

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

        # Download
        print(f"⏳ Downloading {symbol}... (may take 1-2 minutes)")
        data = client.timeseries.get_range(
            dataset=DATASET,
            symbols=[symbol],
            schema=SCHEMA,
            start=start_dt.isoformat(),
            end=end_dt.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: {output_file}")
        print(f"  Size: {file_size:,} bytes ({file_size_kb:.2f} KB)")

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

            print(f"  Bars: {record_count:,}")
            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}")
        except Exception as e:
            print(f"  ⚠️ Could not verify: {e}")

        return output_file, True

    except Exception as e:
        print(f"❌ Download failed: {e}")
        return None, False


def merge_dbn_files(files, output_file):
    """Merge multiple DBN files into one and convert to dataframe."""
    print()
    print("=" * 80)
    print("🔀 MERGING CONTRACTS")
    print("=" * 80)

    all_dfs = []
    for file in files:
        if file and os.path.exists(file):
            print(f"  Loading {file}...")
            store = db.DBNStore.from_file(file)
            df = store.to_df()
            print(f"    Bars: {len(df):,}")
            all_dfs.append(df)

    if not all_dfs:
        print("❌ No data to merge!")
        return None

    # Concatenate and sort by timestamp
    import pandas as pd
    merged_df = pd.concat(all_dfs, ignore_index=False)
    merged_df = merged_df.sort_index()

    print()
    print(f"✅ Merged {len(all_dfs)} contracts")
    print(f"  Total bars: {len(merged_df):,}")
    print(f"  Date range: {merged_df.index[0]} to {merged_df.index[-1]}")

    # Calculate market balance
    bullish_bars = (merged_df['close'] > merged_df['open']).sum()
    bullish_pct = 100 * bullish_bars / len(merged_df)
    trend_pct = 100 * (merged_df['close'].iloc[-1] - merged_df['close'].iloc[0]) / merged_df['close'].iloc[0]

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

    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)")

    return merged_df


def save_to_parquet(df, output_file):
    """Save dataframe directly to Parquet."""
    print()
    print("💾 Saving to Parquet...")

    # Reset index to make timestamp a column
    df_reset = df.reset_index()

    # Save to Parquet
    df_reset.to_parquet(output_file, compression='snappy', index=False)

    file_size = os.path.getsize(output_file)
    file_size_mb = file_size / (1024 * 1024)

    print(f"✅ Saved to {output_file}")
    print(f"  Size: {file_size:,} bytes ({file_size_mb:.2f} MB)")

    return output_file


def main():
    """Download and merge ES futures contracts."""
    print("=" * 80)
    print("ES Futures 90-Day Multi-Contract Download")
    print("=" * 80)
    print()

    # Check API key
    if not API_KEY:
        print("❌ ERROR: DATABENTO_API_KEY not found in environment!")
        sys.exit(1)

    os.makedirs(OUTPUT_DIR, exist_ok=True)
    print(f"📁 Output directory: {OUTPUT_DIR}")
    print(f"📊 Schema: {SCHEMA}")
    print(f"📦 Dataset: {DATASET}")
    print(f"📅 Total range: 2024-08-03 to 2024-11-01 (91 days)")
    print()

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

    # Download each contract
    downloaded_files = []
    for contract in CONTRACTS:
        file_path, success = download_contract(client, contract)
        if success and file_path:
            downloaded_files.append(file_path)

    if not downloaded_files:
        print()
        print("❌ No contracts downloaded successfully!")
        sys.exit(1)

    # Merge contracts
    merged_df = merge_dbn_files(downloaded_files, None)

    if merged_df is None:
        print("❌ Failed to merge contracts!")
        sys.exit(1)

    # Save to Parquet
    output_parquet = os.path.join(OUTPUT_DIR, "ES_FUT_unseen_90d.parquet")
    save_to_parquet(merged_df, output_parquet)

    # Clean up temporary DBN files
    print()
    print("🧹 Cleaning up temporary files...")
    for file in downloaded_files:
        if os.path.exists(file):
            os.remove(file)
            print(f"  Removed {file}")

    # Summary
    print()
    print("=" * 80)
    print("✅ SUCCESS!")
    print("=" * 80)
    print()
    print(f"📊 Final output: {output_parquet}")
    print(f"   Bars: {len(merged_df):,}")
    print(f"   Date range: {merged_df.index[0]} to {merged_df.index[-1]}")
    print()
    print("💰 Estimated cost: ~$9.10 (91 days × $0.10/day)")
    print()
    print("📋 NEXT STEPS:")
    print("1. Validate with DQN evaluation:")
    print("   ./test_dqn_evaluation.sh")
    print()
    print("2. Compare with old 10-day data:")
    print("   # Old: test_data/ES_FUT_unseen.parquet (10 days, Oct 20-30)")
    print("   # New: test_data/ES_FUT_unseen_90d.parquet (91 days, Aug 3 - Nov 1)")
    print()


if __name__ == "__main__":
    main()