#!/usr/bin/env python3
"""
Download 180 days of ZN.FUT (10-Year Treasury Note) data from Databento.
Agent W12-05 - Part of ML Training Roadmap Phase 1.
"""

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

# Configuration
API_KEY = os.getenv("DATABENTO_API_KEY")
OUTPUT_FILE = "/home/jgrusewski/Work/foxhunt/test_data/ZN_FUT_180d.dbn"
LOG_FILE = "/tmp/zn_fut_download_log.txt"
SCHEMA = "ohlcv-1m"
DATASET = "GLBX.MDP3"
# Use continuous contract notation (front month)
SYMBOL = "ZN.c.0"

# Date range: 180 days ending 2024-10-20 (2025 data not yet available)
START_DATE = "2024-04-23"
END_DATE = "2024-10-20"

def log(message):
    """Write to both console and log file."""
    print(message)
    with open(LOG_FILE, "a") as f:
        f.write(message + "\n")

def main():
    """Download ZN.FUT 180-day data."""
    log("=" * 80)
    log("Agent W12-05: Download ZN.FUT (180 days)")
    log("=" * 80)
    log(f"Symbol: {SYMBOL} (10-Year Treasury Note)")
    log(f"Date Range: {START_DATE} to {END_DATE} (180 days)")
    log(f"Schema: {SCHEMA}")
    log(f"Dataset: {DATASET}")
    log(f"Output: {OUTPUT_FILE}")
    log("")

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

    log(f"✅ API key found (length: {len(API_KEY)})")

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

    # Download data
    log("")
    log(f"📥 Downloading {SYMBOL} data...")
    log(f"   Start: {START_DATE}")
    log(f"   End:   {END_DATE}")

    try:
        # Parse dates
        start_dt = datetime.strptime(START_DATE, "%Y-%m-%d").replace(tzinfo=timezone.utc)
        # Use end of day for historical 2024 data
        end_dt = datetime.strptime(END_DATE, "%Y-%m-%d").replace(hour=23, minute=59, second=59, tzinfo=timezone.utc)

        # Request data
        log("   Requesting data from Databento API...")
        data = client.timeseries.get_range(
            dataset=DATASET,
            symbols=[SYMBOL],
            schema=SCHEMA,
            start=start_dt.isoformat(),
            end=end_dt.isoformat(),
        )

        # Create output directory if needed
        os.makedirs(os.path.dirname(OUTPUT_FILE), exist_ok=True)

        # Write to file
        log("   Writing to file...")
        data.to_file(OUTPUT_FILE)

        # Verify data
        file_size = os.path.getsize(OUTPUT_FILE)
        file_size_mb = file_size / (1024 * 1024)
        log(f"✅ File created: {OUTPUT_FILE}")
        log(f"   Size: {file_size_mb:.2f} MB ({file_size:,} bytes)")

        # Parse and verify bars
        try:
            store = db.DBNStore.from_file(OUTPUT_FILE)
            df = store.to_df()
            bar_count = len(df)

            log(f"   Bars: {bar_count:,}")

            if bar_count > 0:
                log(f"   Price Range: {float(df['close'].min()):.2f} - {float(df['close'].max()):.2f}")
                log(f"   Total Volume: {int(df['volume'].sum()):,}")
                log(f"   Date Range: {df.index[0]} to {df.index[-1]}")

            # Cost estimation (rough: $0.003-0.004 per bar)
            estimated_cost = bar_count * 0.0035 / 1000
            log(f"   Estimated Cost: ${estimated_cost:.2f}")

            log("")
            log("=" * 80)
            log("✅ ZN.FUT DOWNLOAD COMPLETE")
            log("=" * 80)
            log(f"File: {OUTPUT_FILE}")
            log(f"Size: {file_size_mb:.2f} MB")
            log(f"Bars: {bar_count:,}")
            log(f"Cost: ~${estimated_cost:.2f}")
            log("")

            # Success criteria check
            if 800_000 <= bar_count <= 1_000_000:
                log("✅ SUCCESS: Bar count within expected range (800K-1.0M)")
            else:
                log(f"⚠️  WARNING: Bar count {bar_count:,} outside expected range (800K-1.0M)")

            if 55 <= file_size_mb <= 75:
                log("✅ SUCCESS: File size within expected range (55-75 MB)")
            else:
                log(f"⚠️  WARNING: File size {file_size_mb:.2f} MB outside expected range (55-75 MB)")

            if estimated_cost <= 0.80:
                log("✅ SUCCESS: Cost within budget (≤$0.80)")
            else:
                log(f"⚠️  WARNING: Cost ${estimated_cost:.2f} exceeds budget ($0.80)")

            # Write completion flag
            with open("/tmp/w12_05_complete.flag", "w") as f:
                f.write(f"ZN.FUT download complete\n")
                f.write(f"File: {OUTPUT_FILE}\n")
                f.write(f"Size: {file_size_mb:.2f} MB\n")
                f.write(f"Bars: {bar_count:,}\n")
                f.write(f"Cost: ${estimated_cost:.2f}\n")

            return 0

        except Exception as e:
            log(f"⚠️  WARNING: Could not verify file: {e}")
            log(f"   File created but verification failed")
            return 1

    except Exception as e:
        log("")
        log(f"❌ ERROR: Download failed: {e}")
        log("")
        import traceback
        log(traceback.format_exc())
        sys.exit(1)

if __name__ == "__main__":
    # Clear log file
    with open(LOG_FILE, "w") as f:
        f.write("")

    sys.exit(main())