#!/usr/bin/env python3
"""
Hypersets Demo

This demo showcases all the key features of Hypersets:
• Fast metadata retrieval without downloading entire datasets
• Efficient SQL querying with DuckDB optimization
• Smart caching to avoid repeated API calls
• Optional download tracking with data savings estimation
• Multiple output formats (pandas, HuggingFace Dataset)
• Proper error handling and rate limiting
• Memory-only operation

We test with real datasets to prove the implementation works.
"""

import logging
import time
import hypersets as hs

# Configure logging to see what's happening
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)


def test_dataset_info_retrieval():
    """Test 1: Dataset info retrieval with caching"""
    print("🔍 Test 1: Dataset Info Retrieval")
    print("=" * 50)

    print("Testing with OMARKAMALI/WIKIPEDIA-MONTHLY dataset...")
    start = time.time()

    # Get dataset info
    info = hs.get_dataset_info("omarkamali/wikipedia-monthly")

    elapsed = time.time() - start
    print(f"✅ Retrieved dataset info in {elapsed:.2f}s")
    print(f"   Dataset: {info.name}")
    print(f"   Configs: {info.config_names[:50]}...")  # Show first 50 configs
    print(f"   Total files: {info.total_parquet_files}")
    print(f"   Estimated size: {info.estimated_total_size_gb:.1f} GB")

    # Test caching by calling again
    print(f"   Available configs: {len(info.config_names)}")

    print("\nTesting with Wikipedia (complex dataset)...")
    start = time.time()

    # Test with different dataset
    wiki_info = hs.get_dataset_info("wikimedia/wikipedia")

    elapsed = time.time() - start
    print(f"✅ Retrieved Wikipedia info in {elapsed:.2f}s")
    print(f"   Total configs: {len(wiki_info.config_names)}")
    print(f"   Example configs: {wiki_info.config_names[:5]}")
    print(f"   Total files: {wiki_info.total_parquet_files}")
    print(f"   Estimated size: {wiki_info.estimated_total_size_gb:.1f} GB")


def test_efficient_querying():
    """Test 2: Efficient SQL querying with DuckDB"""
    print(f"\n{'🚀 Test 2: Efficient Querying'}")
    print("=" * 50)

    print("Testing basic SQL query...")
    start = time.time()

    # Simple query with minimal memory usage
    result = hs.query(
        "SELECT title, LENGTH(text) as text_length FROM dataset LIMIT 5",
        dataset="omarkamali/wikipedia-monthly",
        config="latest.en",
    )

    elapsed = time.time() - start
    print(f"✅ Query completed in {elapsed:.2f}s")
    print(f"   Rows returned: {len(result)}")
    print(f"   Columns: {result.columns}")
    print(f"   Shape: {result.shape}")

    df = result.to_pandas()
    print("\n   Sample results:")
    for i, row in df.head(3).iterrows():
        title = row["title"][:50] + "..." if len(row["title"]) > 50 else row["title"]
        print(f"     {i+1}. {title} ({row['text_length']} chars)")

    print("\nTesting aggregation query...")
    start = time.time()

    # Fast aggregation query
    count_result = hs.query(
        "SELECT COUNT(*) as count FROM dataset",
        dataset="omarkamali/wikipedia-monthly",
        config="latest.en",
    )

    elapsed = time.time() - start
    print(
        f"✅ Aggregation completed in {elapsed:.2f}s, result: {count_result.to_pandas()}"
    )


def test_sampling_capabilities():
    """Test 3: Sampling with DuckDB optimization"""
    print(f"\n{'🎲 Test 3: Sampling Capabilities'}")
    print("=" * 50)

    print("Testing random sampling...")
    start = time.time()

    # Use DuckDB's native sampling with minimal columns
    sample_result = hs.sample(
        n=10,
        dataset="omarkamali/wikipedia-monthly",
        columns=["title", "LENGTH(text) as text_length"],
        config="latest.en",
    )

    elapsed = time.time() - start
    print(f"✅ Sampled {len(sample_result)} rows in {elapsed:.2f}s")

    df = sample_result.to_pandas()

    print("\n   Sample articles:")
    for i, row in df.head(5).iterrows():
        title = row["title"][:60] + "..." if len(row["title"]) > 60 else row["title"]
        print(f"     {i+1}. {title} ({row['text_length']} chars)")


def test_download_tracking():
    """Test 4: Download tracking and data savings"""
    print(f"\n{'📊 Test 4: Download Tracking & Data Savings'}")
    print("=" * 50)

    print("Testing query with download tracking...")
    start = time.time()

    # Query with download tracking enabled
    result = hs.query(
        "SELECT title, url FROM dataset LIMIT 100",
        dataset="omarkamali/wikipedia-monthly",
        config="latest.en",
        track_downloads=True,
    )

    elapsed = time.time() - start
    print(f"✅ Query completed")
    print(f"   Rows: {len(result)}")
    print(f"   Query time: {elapsed:.2f}s")

    # Show download tracking info
    if hasattr(result, "download_stats") and result.download_stats:
        stats = result.download_stats
        print(f"\n💰 Data Savings Analysis:")
        print(f"   Total dataset size: {stats.total_dataset_size_gb:.1f} GB")
        print(f"   Estimated downloaded: {stats.estimated_downloaded_gb:.2f} GB")
        print(
            f"   Savings: {stats.data_saved_gb:.1f} GB ({stats.savings_percentage:.1f}%)"
        )
        print(f"   {stats.savings_message}")


def test_output_formats():
    """Test 5: Multiple output formats"""
    print(f"\n{'📋 Test 5: Output Formats'}")
    print("=" * 50)

    # Test pandas output (memory-efficient)
    result = hs.head(
        n=5,
        dataset="omarkamali/wikipedia-monthly",
        config="latest.en",
        columns=["title", "url"],
    )

    # Test pandas output
    df = result.to_pandas()
    print(f"✅ Pandas DataFrame: {df.shape}")
    print(f"   Columns: {list(df.columns)}")
    print(f"   Sample: {df.iloc[0]['title'][:50]}...")

    # Test HuggingFace Dataset output
    try:
        hf_dataset = result.to_dataset()
        print(f"✅ HuggingFace Dataset: {len(hf_dataset)} rows")
        print(f"   Features: {list(hf_dataset.features.keys())}")
    except Exception as e:
        print(f"⚠️  HuggingFace conversion: {str(e)}")


def test_schema_inspection():
    """Test 6: Schema inspection"""
    print(f"\n{'🔍 Test 6: Schema Inspection'}")
    print("=" * 50)

    # Get schema information
    schema = hs.schema(dataset="omarkamali/wikipedia-monthly", config="latest.en")

    print(f"✅ Schema retrieved:")
    print(f"   Columns: {len(schema)}")
    for i, row in schema.head(5).iterrows():  # Show first 5 columns
        print(f"     • {row['column_name']}: {row['column_type']}")
    if len(schema) > 5:
        print(f"     ... and {len(schema) - 5} more columns")


def test_caching_performance():
    """Test 7: Caching performance"""
    print(f"\n{'⚡ Test 7: Caching Performance'}")
    print("=" * 50)

    # Clear cache and time first call
    hs.clear_cache()
    start = time.time()
    info1 = hs.get_dataset_info("omarkamali/wikipedia-monthly")
    first_call = time.time() - start

    # Time second call (should be cached)
    start = time.time()
    info2 = hs.get_dataset_info("omarkamali/wikipedia-monthly")
    second_call = time.time() - start

    speedup = first_call / second_call if second_call > 0 else float("inf")

    print(f"✅ Caching Analysis:")
    print(f"   First call (cold): {first_call:.2f}s")
    print(f"   Second call (cached): {second_call:.3f}s")
    print(f"   Speedup: {speedup:.0f}x")
    print(f"   Cache is working: {speedup > 10}")


def run_demo():
    """Run the complete Hypersets demo"""
    print("🎯 HYPERSETS DEMO")
    print("=" * 60)
    print()
    print("Hypersets enables efficient querying of massive datasets")
    print("without downloading them entirely. Let's see it in action!")
    print()

    try:
        test_dataset_info_retrieval()
        test_efficient_querying()
        test_sampling_capabilities()
        test_download_tracking()
        test_output_formats()
        test_schema_inspection()
        test_caching_performance()

        print(f"\n{'🎉 DEMO COMPLETED SUCCESSFULLY!'}")
    except Exception as e:
        print(f"\n❌ Demo failed with error: {e}")
        import traceback

        traceback.print_exc()


if __name__ == "__main__":
    run_demo()
