#!/usr/bin/env python3
"""
Advanced querying examples for Hypersets.

This script demonstrates sophisticated SQL queries and advanced use cases
for efficient dataset analysis and exploration using Wikipedia data.
"""

import hypersets as hs
import pandas as pd
import time


def measure_time(func):
    """Decorator to measure execution time."""

    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(f"⏱️  Execution time: {end - start:.2f} seconds")
        return result

    return wrapper


@measure_time
def text_length_analysis():
    """Analyze text length distribution in Wikipedia articles."""
    print("📊 Text Length Analysis")
    print("-" * 40)

    query = """
    WITH text_stats AS (
        SELECT 
            LENGTH(text) as text_length,
            LENGTH(title) as title_length,
            CASE 
                WHEN LENGTH(text) < 1000 THEN 'Very Short'
                WHEN LENGTH(text) < 5000 THEN 'Short'
                WHEN LENGTH(text) < 20000 THEN 'Medium'
                WHEN LENGTH(text) < 50000 THEN 'Long'
                ELSE 'Very Long'
            END as length_category
        FROM dataset
        TABLESAMPLE(1000 ROWS)
    )
    SELECT 
        length_category,
        COUNT(*) as count,
        AVG(text_length) as avg_text_length,
        AVG(title_length) as avg_title_length,
        MIN(text_length) as min_text_length,
        MAX(text_length) as max_text_length
    FROM text_stats
    GROUP BY length_category
    ORDER BY avg_text_length
    """

    try:
        result = hs.query(
            query, dataset="omarkamali/wikipedia-monthly", config="latest.en"
        )

        df = result.to_pandas()
        print("Text Length Distribution:")
        print(df.to_string(index=False))

        return df
    except Exception as e:
        print(f"Error: {e}")


@measure_time
def keyword_search_analysis():
    """Search for articles containing specific keywords."""
    print("\n🔍 Keyword Search Analysis")
    print("-" * 40)

    query = """
    SELECT 
        title,
        LENGTH(text) as article_length,
        CASE 
            WHEN text ILIKE '%artificial intelligence%' THEN 'AI'
            WHEN text ILIKE '%machine learning%' THEN 'ML'
            WHEN text ILIKE '%deep learning%' THEN 'Deep Learning'
            WHEN text ILIKE '%neural network%' THEN 'Neural Networks'
            ELSE 'Other Tech'
        END as tech_category
    FROM dataset
    WHERE text ILIKE '%artificial intelligence%' 
       OR text ILIKE '%machine learning%'
       OR text ILIKE '%deep learning%'
       OR text ILIKE '%neural network%'
    LIMIT 20
    """

    try:
        result = hs.query(
            query, dataset="omarkamali/wikipedia-monthly", config="latest.en"
        )

        df = result.to_pandas()
        print(f"Found {len(df)} tech-related articles:")

        # Group by category
        category_counts = df["tech_category"].value_counts()
        print("\nTech Category Distribution:")
        for category, count in category_counts.items():
            print(f"  {category}: {count}")

        print("\nSample articles:")
        for i, row in df.head(5).iterrows():
            print(f"  • {row['title'][:60]}... ({row['article_length']:,} chars)")

        return df
    except Exception as e:
        print(f"Error: {e}")


@measure_time
def url_pattern_analysis():
    """Analyze URL patterns and domains."""
    print("\n🌐 URL Pattern Analysis")
    print("-" * 40)

    query = """
    SELECT 
        SPLIT_PART(url, '/', 3) as domain,
        COUNT(*) as article_count,
        AVG(LENGTH(text)) as avg_text_length
    FROM dataset
    WHERE url IS NOT NULL AND url != ''
    GROUP BY SPLIT_PART(url, '/', 3)
    ORDER BY article_count DESC
    LIMIT 10
    """

    try:
        result = hs.query(
            query, dataset="omarkamali/wikipedia-monthly", config="latest.en"
        )

        df = result.to_pandas()
        print("Top domains by article count:")
        print(df.to_string(index=False))

        return df
    except Exception as e:
        print(f"Error: {e}")


@measure_time
def title_pattern_analysis():
    """Analyze Wikipedia title patterns."""
    print("\n📝 Title Pattern Analysis")
    print("-" * 40)

    query = """
    WITH title_analysis AS (
        SELECT 
            title,
            LENGTH(title) as title_length,
            LENGTH(text) as text_length,
            CASE 
                WHEN title LIKE '%(%' THEN 'Has Parentheses'
                WHEN title LIKE '%:%' THEN 'Has Colon'
                WHEN title LIKE 'List of%' THEN 'List Article'
                WHEN title LIKE 'History of%' THEN 'History Article'
                ELSE 'Standard Title'
            END as title_type
        FROM dataset
        TABLESAMPLE(2000 ROWS)
    )
    SELECT 
        title_type,
        COUNT(*) as count,
        AVG(title_length) as avg_title_length,
        AVG(text_length) as avg_text_length,
        COUNT(*) * 100.0 / SUM(COUNT(*)) OVER() as percentage
    FROM title_analysis
    GROUP BY title_type
    ORDER BY count DESC
    """

    try:
        result = hs.query(
            query, dataset="omarkamali/wikipedia-monthly", config="latest.en"
        )

        df = result.to_pandas()
        print("Title Pattern Analysis:")
        print(df.to_string(index=False, float_format="%.1f"))

        return df
    except Exception as e:
        print(f"Error: {e}")


@measure_time
def content_quality_metrics():
    """Calculate content quality metrics."""
    print("\n📈 Content Quality Metrics")
    print("-" * 40)

    query = """
    WITH quality_metrics AS (
        SELECT 
            title,
            LENGTH(text) as text_length,
            LENGTH(title) as title_length,
            REGEXP_COUNT(text, '\\[\\[([^\\]]+)\\]\\]') as internal_links,
            REGEXP_COUNT(text, 'http[s]?://[^\\s]+') as external_links,
            REGEXP_COUNT(text, '\\n== [^=]+ ==\\n') as section_count,
            CASE 
                WHEN LENGTH(text) > 10000 AND REGEXP_COUNT(text, '\\[\\[([^\\]]+)\\]\\]') > 10 THEN 'High Quality'
                WHEN LENGTH(text) > 5000 AND REGEXP_COUNT(text, '\\[\\[([^\\]]+)\\]\\]') > 5 THEN 'Medium Quality'
                ELSE 'Basic Quality'
            END as quality_tier
        FROM dataset
        TABLESAMPLE(1000 ROWS)
    )
    SELECT 
        quality_tier,
        COUNT(*) as article_count,
        AVG(text_length) as avg_text_length,
        AVG(internal_links) as avg_internal_links,
        AVG(external_links) as avg_external_links,
        AVG(section_count) as avg_sections
    FROM quality_metrics
    GROUP BY quality_tier
    ORDER BY avg_text_length DESC
    """

    try:
        result = hs.query(
            query, dataset="omarkamali/wikipedia-monthly", config="latest.en"
        )

        df = result.to_pandas()
        print("Content Quality Analysis:")
        print(df.to_string(index=False, float_format="%.1f"))

        return df
    except Exception as e:
        print(f"Error: {e}")


@measure_time
def performance_comparison():
    """Compare different query patterns for performance."""
    print("\n⚡ Performance Comparison")
    print("-" * 40)

    # Test 1: Simple count
    print("Test 1: Simple COUNT query")
    start = time.time()
    count_result = hs.count(dataset="omarkamali/wikipedia-monthly", config="latest.en")
    count_time = time.time() - start
    print(f"  Total articles: {count_result:,}")
    print(f"  Time: {count_time:.2f}s")

    # Test 2: Aggregation with sampling
    print("\nTest 2: Complex aggregation with sampling")
    start = time.time()
    sample_result = hs.query(
        """
        SELECT 
            COUNT(*) as total_articles,
            AVG(LENGTH(text)) as avg_length,
            MAX(LENGTH(text)) as max_length
        FROM (
            SELECT text 
            FROM dataset 
            TABLESAMPLE(5000 ROWS)
        )
        """,
        dataset="omarkamali/wikipedia-monthly",
        config="latest.en",
    )
    agg_time = time.time() - start
    stats = sample_result.to_pandas().iloc[0]
    print(f"  Sample size: {stats['total_articles']:,}")
    print(f"  Avg length: {stats['avg_length']:.0f} chars")
    print(f"  Max length: {stats['max_length']:,} chars")
    print(f"  Time: {agg_time:.2f}s")

    # Test 3: Text search
    print("\nTest 3: Text search query")
    start = time.time()
    search_result = hs.query(
        """
        SELECT COUNT(*) as matches
        FROM dataset
        WHERE text ILIKE '%computer%'
        LIMIT 1000
        """,
        dataset="omarkamali/wikipedia-monthly",
        config="latest.en",
    )
    search_time = time.time() - start
    matches = search_result.to_pandas().iloc[0]["matches"]
    print(f"  Matches found: {matches:,}")
    print(f"  Time: {search_time:.2f}s")


def main():
    """Run all advanced query examples."""
    print("🚀 Hypersets Advanced Query Examples")
    print("=" * 60)
    print("Using Wikipedia Monthly dataset for analysis\n")

    try:
        # Get dataset info first
        print("📋 Dataset Information")
        print("-" * 40)
        info = hs.get_dataset_info("omarkamali/wikipedia-monthly")
        print(f"Dataset: {info.name}")
        print(f"Available configs: {len(info.config_names)}")
        print(f"Total files: {info.total_parquet_files}")
        print(f"Estimated size: {info.estimated_total_size_gb:.1f} GB")

        if "latest.en" in info.configs:
            en_config = info.configs["latest.en"]
            train_split = en_config.splits.get("train")
            if train_split:
                print(f"English articles: {train_split.file_count} files")
        print()

        # Run analysis functions
        text_length_analysis()
        keyword_search_analysis()
        url_pattern_analysis()
        title_pattern_analysis()
        content_quality_metrics()
        performance_comparison()

        print("\n" + "=" * 60)
        print("✅ All advanced query examples completed!")
        print("=" * 60)
        print("\n💡 Key learnings:")
        print("  • Use TABLESAMPLE for faster exploration of large datasets")
        print("  • Combine aggregations with window functions for insights")
        print("  • Regular expressions enable powerful text analysis")
        print("  • DuckDB's SQL dialect supports advanced analytics")
        print("  • Hypersets makes TB-scale analysis accessible")
        print("\n🚀 Try creating your own advanced queries!")

    except Exception as e:
        print(f"❌ Error running examples: {e}")
        import traceback

        traceback.print_exc()


if __name__ == "__main__":
    main()
