"""
Benchmark Dataset Generator.

Generates synthetic CSV datasets of varying sizes for stress testing.
Includes realistic data patterns:
- Categorical data (low/high cardinality)
- Numerical data (ints, floats)
- Dates and timestamps
- Text fields
- Missing values (nulls)
- Duplicate rows
"""

import polars as pl
import numpy as np
import os
import time
from pathlib import Path

def generate_dataset(rows: int, output_path: str, seed: int = 42):
    """Generate a synthetic dataset with Polars."""
    print(f"Generating {rows:,} rows to {output_path}...")
    
    rng = np.random.default_rng(seed)
    
    # Define column data
    data = {
        "id": np.arange(rows),
        "category_low": rng.choice(["A", "B", "C", "D", "E"], rows),
        "category_high": rng.choice([f"CAT_{i}" for i in range(100)], rows),
        "value_int": rng.integers(0, 100000, rows),
        "value_float": rng.random(rows) * 1000,
        "date": rng.choice(
            pl.date_range(
                start=pl.date(2020, 1, 1),
                end=pl.date(2023, 12, 31),
                interval="1d",
                eager=True
            ),
            rows
        ),
        "text_short": rng.choice(["active", "inactive", "pending", "archived"], rows),
        "text_long": rng.choice(
            [
                "The quick brown fox jumps over the lazy dog",
                "Lorem ipsum dolor sit amet",
                "Data cleaning is essential for ML pipelines",
                "Polars is much faster than Pandas"
            ], 
            rows
        )
    }
    
    df = pl.DataFrame(data)
    
    # Introduce nulls (5% random nulls in some columns)
    df = df.with_columns([
        pl.when(pl.lit(rng.random(rows) < 0.05))
        .then(None)
        .otherwise(pl.col("value_float"))
        .alias("value_float"),
        
        pl.when(pl.lit(rng.random(rows) < 0.05))
        .then(None)
        .otherwise(pl.col("text_short"))
        .alias("text_short")
    ])
    
    # Introduce duplicates (append 1% of rows again)
    dupes = df.sample(fraction=0.01, seed=seed)
    df = pl.concat([df, dupes])
    
    # Save to CSV
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    df.write_csv(output_path)
    
    size_mb = os.path.getsize(output_path) / (1024 * 1024)
    print(f"✅ Generated {output_path} ({size_mb:.2f} MB)")
    return size_mb

def main():
    """Generate Small, Medium, and Large datasets."""
    base_dir = Path("benchmarks/data")
    base_dir.mkdir(parents=True, exist_ok=True)
    
    # 1. Small (100k rows, ~10MB)
    generate_dataset(100_000, str(base_dir / "dataset_small.csv"))
    
    # 2. Medium (1M rows, ~100MB)
    generate_dataset(1_000_000, str(base_dir / "dataset_medium.csv"))
    
    # 3. Large (10M rows, ~1GB)
    # Uncomment for full stress test (takes longer)
    # generate_dataset(10_000_000, str(base_dir / "dataset_large.csv"))
    
    print("\n⚠️  Note: Large dataset (1GB) generation commented out for speed.")
    print("Run with 10M rows to generate full 1GB file.")

if __name__ == "__main__":
    main()
