#!/usr/bin/env python3
"""
Download 180 days of ZN.FUT (10-Year Treasury Note) data from Databento.
Agent W12-05 - Downloads day-by-day and concatenates.
"""

import os
import sys
from datetime import datetime, timedelta, timezone
import databento as db
import tempfile
import shutil

# 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"
SYMBOL = "ZN.FUT"

# 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, end="\n"):
    """Write to both console and log file."""
    print(message, end=end, flush=True)
    with open(LOG_FILE, "a") as f:
        f.write(message + (end if end != "\n" else "\n"))

def generate_trading_days(start_str, end_str):
    """Generate list of trading days (excluding weekends)."""
    start = datetime.strptime(start_str, "%Y-%m-%d")
    end = datetime.strptime(end_str, "%Y-%m-%d")

    dates = []
    current = start
    while current <= end:
        # Skip weekends (Saturday=5, Sunday=6)
        if current.weekday() < 5:
            dates.append(current.strftime("%Y-%m-%d"))
        current += timedelta(days=1)

    return dates

def download_day(client, date_str, temp_dir):
    """Download data for a single day."""
    try:
        # Parse date (full trading day UTC)
        start_dt = datetime.strptime(date_str, "%Y-%m-%d").replace(tzinfo=timezone.utc)
        end_dt = start_dt.replace(hour=23, minute=59, second=59)

        # Output to temp file
        temp_file = os.path.join(temp_dir, f"ZN_{date_str}.dbn")

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

        # Write to temp file
        data.to_file(temp_file)

        # Check file size
        file_size = os.path.getsize(temp_file)

        if file_size < 1000:  # Less than 1KB = no data (holiday)
            os.remove(temp_file)
            return {"status": "skipped", "reason": "no data (holiday)", "bars": 0}

        # Parse and count bars
        try:
            store = db.DBNStore.from_file(temp_file)
            df = store.to_df()
            bar_count = len(df)
            return {"status": "success", "bars": bar_count, "file": temp_file}
        except:
            return {"status": "success", "bars": "?", "file": temp_file}

    except Exception as e:
        return {"status": "error", "reason": str(e)}

def concatenate_files(file_list, output_file):
    """Concatenate multiple DBN files into one."""
    log(f"\n📦 Concatenating {len(file_list)} files...")

    # Read all files and combine
    all_data = []
    for file_path in file_list:
        store = db.DBNStore.from_file(file_path)
        all_data.append(store)

    # Write combined data
    # For now, just copy the first file and append the rest
    # (DBN doesn't have a built-in concat, so we'll use file copy)
    with open(output_file, 'wb') as outf:
        for file_path in file_list:
            with open(file_path, 'rb') as inf:
                outf.write(inf.read())

    log(f"✅ Concatenation complete")

def main():
    """Download ZN.FUT 180-day data day-by-day."""
    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)

    # Generate trading days
    trading_days = generate_trading_days(START_DATE, END_DATE)
    log(f"📅 Trading days to download: {len(trading_days)}")
    log("")

    # Create temp directory
    temp_dir = tempfile.mkdtemp(prefix="zn_download_")
    log(f"📂 Temp directory: {temp_dir}")
    log("")

    # Download day by day
    log("📥 Downloading daily data...")
    successful_files = []
    total_bars = 0
    skipped = 0
    errors = 0

    for i, date_str in enumerate(trading_days):
        progress = (i + 1) / len(trading_days) * 100
        log(f"[{i+1}/{len(trading_days)} - {progress:.1f}%] {date_str}...", end=" ")

        result = download_day(client, date_str, temp_dir)

        if result["status"] == "success":
            bars = result["bars"]
            total_bars += bars if isinstance(bars, int) else 0
            successful_files.append(result["file"])
            log(f"✅ {bars} bars")
        elif result["status"] == "skipped":
            skipped += 1
            log(f"⏭️  {result['reason']}")
        elif result["status"] == "error":
            errors += 1
            log(f"❌ {result['reason']}")

    log("")
    log(f"✅ Downloaded: {len(successful_files)} days")
    log(f"⏭️  Skipped: {skipped} days (holidays)")
    log(f"❌ Errors: {errors} days")
    log(f"📊 Total bars: {total_bars:,}")
    log("")

    if not successful_files:
        log("❌ ERROR: No data downloaded!")
        shutil.rmtree(temp_dir)
        sys.exit(1)

    # Concatenate files
    concatenate_files(successful_files, OUTPUT_FILE)

    # Clean up temp directory
    shutil.rmtree(temp_dir)
    log(f"🗑️  Cleaned up temp directory")
    log("")

    # Verify final file
    file_size = os.path.getsize(OUTPUT_FILE)
    file_size_mb = file_size / (1024 * 1024)

    log("=" * 80)
    log("✅ ZN.FUT DOWNLOAD COMPLETE")
    log("=" * 80)
    log(f"File: {OUTPUT_FILE}")
    log(f"Size: {file_size_mb:.2f} MB ({file_size:,} bytes)")
    log(f"Bars: {total_bars:,}")

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

    # Success criteria check
    if 800_000 <= total_bars <= 1_000_000:
        log("✅ SUCCESS: Bar count within expected range (800K-1.0M)")
    else:
        log(f"⚠️  WARNING: Bar count {total_bars:,} 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: {total_bars:,}\n")
        f.write(f"Cost: ${estimated_cost:.2f}\n")

    return 0

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

    sys.exit(main())