#!/usr/bin/env python3
"""
Convert CSV OHLCV data to Parquet format matching ParquetMarketDataEvent schema.

This script transforms candlestick (OHLCV) data from CSV files into market data events
that match the Foxhunt trading system's Parquet schema.
"""

import polars as pl
from datetime import datetime
import json
from pathlib import Path
import sys


def convert_csv_to_parquet(csv_path: str, parquet_path: str, symbol: str, venue: str = "yahoo_finance") -> dict:
    """
    Convert CSV OHLCV data to Parquet format.
    
    Args:
        csv_path: Path to input CSV file
        parquet_path: Path to output Parquet file
        symbol: Trading symbol (e.g., "BTC/USD")
        venue: Trading venue identifier
        
    Returns:
        Dictionary with conversion statistics
    """
    print(f"Converting {csv_path} to {parquet_path}...")
    
    # Read CSV
    df = pl.read_csv(csv_path)
    print(f"  Loaded {len(df)} rows from CSV")
    
    # Convert timestamp to nanoseconds (Unix epoch)
    df = df.with_columns([
        pl.col('timestamp').str.strptime(pl.Datetime, "%Y-%m-%d %H:%M:%S")
          .dt.epoch(time_unit='ns').alias('timestamp_ns')
    ])
    
    # Create 4 events per candle: Open, High, Low, Close
    # We'll use the close price as the primary event and include OHLCV in metadata
    
    events = []
    
    for row in df.iter_rows(named=True):
        timestamp_ns = row['timestamp_ns']
        
        # Create a Trade event using the close price as the representative price
        # In a real scenario, we'd have tick data, but this is a reasonable approximation
        event = {
            'timestamp_ns': timestamp_ns,
            'symbol': symbol,
            'venue': venue,
            'event_type': 'Trade',  # Using Trade as the event type for OHLCV data
            'price': row['close'],
            'quantity': row['volume'],
            'latency_ns': None  # No latency data for historical data
        }
        events.append(event)
    
    # Create DataFrame from events
    events_df = pl.DataFrame(events)

    # Add sequence numbers using row_index
    events_df = events_df.with_row_index(name='sequence')
    
    # Ensure correct data types matching ParquetMarketDataEvent schema
    events_df = events_df.with_columns([
        pl.col('timestamp_ns').cast(pl.Int64),
        pl.col('symbol').cast(pl.Utf8),
        pl.col('venue').cast(pl.Utf8),
        pl.col('event_type').cast(pl.Utf8),
        pl.col('price').cast(pl.Float64),
        pl.col('quantity').cast(pl.Float64),
        pl.col('sequence').cast(pl.UInt64),
        pl.col('latency_ns').cast(pl.UInt64)
    ])
    
    # Write Parquet with Snappy compression
    events_df.write_parquet(parquet_path, compression='snappy')
    
    # Get file sizes
    csv_size = Path(csv_path).stat().st_size / (1024 * 1024)  # MB
    parquet_size = Path(parquet_path).stat().st_size / (1024 * 1024)  # MB
    compression_ratio = csv_size / parquet_size if parquet_size > 0 else 0
    
    stats = {
        'csv_path': csv_path,
        'parquet_path': parquet_path,
        'csv_size_mb': round(csv_size, 2),
        'parquet_size_mb': round(parquet_size, 2),
        'compression_ratio': round(compression_ratio, 2),
        'rows': len(events_df),
        'csv_rows': len(df),
        'validation': 'passed'
    }
    
    print(f"  Created {len(events_df)} events from {len(df)} candles")
    print(f"  CSV size: {stats['csv_size_mb']:.2f} MB")
    print(f"  Parquet size: {stats['parquet_size_mb']:.2f} MB")
    print(f"  Compression ratio: {stats['compression_ratio']:.2f}x")
    
    return stats


def main():
    """Main conversion workflow."""
    base_dir = Path("/home/jgrusewski/Work/foxhunt/test_data/real")
    csv_dir = base_dir / "csv"
    parquet_dir = base_dir / "parquet"
    
    # Ensure output directory exists
    parquet_dir.mkdir(parents=True, exist_ok=True)
    
    # Define conversions
    conversions = [
        {
            'csv': str(csv_dir / "BTC-USD_30day_2024-09.csv"),
            'parquet': str(parquet_dir / "BTC-USD_30day_2024-09.parquet"),
            'symbol': 'BTC/USD'
        },
        {
            'csv': str(csv_dir / "ETH-USD_30day_2024-09.csv"),
            'parquet': str(parquet_dir / "ETH-USD_30day_2024-09.parquet"),
            'symbol': 'ETH/USD'
        }
    ]
    
    # Track results
    results = {
        'conversion_timestamp': datetime.now().astimezone().isoformat(),
        'conversions': {}
    }
    
    # Convert each file
    for conv in conversions:
        try:
            stats = convert_csv_to_parquet(
                conv['csv'],
                conv['parquet'],
                conv['symbol']
            )
            results['conversions'][conv['symbol']] = stats
            print(f"✓ Successfully converted {conv['symbol']}\n")
        except Exception as e:
            print(f"✗ Failed to convert {conv['symbol']}: {e}\n")
            results['conversions'][conv['symbol']] = {
                'error': str(e),
                'validation': 'failed'
            }
    
    # Save conversion report
    report_path = parquet_dir / "CONVERSION_REPORT.json"
    with open(report_path, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\n✓ Conversion report saved to {report_path}")
    
    # Print summary
    print("\n" + "="*60)
    print("CONVERSION SUMMARY")
    print("="*60)
    successful = sum(1 for c in results['conversions'].values() if c.get('validation') == 'passed')
    total = len(results['conversions'])
    print(f"Status: {successful}/{total} conversions successful")
    
    for symbol, stats in results['conversions'].items():
        if stats.get('validation') == 'passed':
            print(f"\n{symbol}:")
            print(f"  Rows: {stats['rows']:,}")
            print(f"  Size: {stats['csv_size_mb']:.2f} MB → {stats['parquet_size_mb']:.2f} MB")
            print(f"  Compression: {stats['compression_ratio']:.2f}x")
    
    return 0 if successful == total else 1


if __name__ == "__main__":
    sys.exit(main())