"""
Stress Test Suite.

Benchmarks the PolarsEngine against generated datasets.
Measures:
- Execution time (Load, Clean, Save)
- Memory usage (Peak)
- Throughput (Rows/sec)
"""

import time
import os
import psutil
import polars as pl
import sys
from pathlib import Path

# Add backend to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from backend.app.engines.polars_engine import PolarsEngine

def get_memory_usage():
    """Get current memory usage in MB."""
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / (1024 * 1024)

def run_benchmark(file_path: str):
    """Run benchmark on a single file."""
    print(f"\n🚀 Benchmarking: {file_path}")
    
    if not os.path.exists(file_path):
        print(f"❌ File not found: {file_path}")
        return
    
    file_size_mb = os.path.getsize(file_path) / (1024 * 1024)
    print(f"Size: {file_size_mb:.2f} MB")
    
    engine = PolarsEngine()
    
    # Measure Load Time
    start_mem = get_memory_usage()
    start_time = time.time()
    
    df = pl.read_csv(file_path)
    
    load_time = time.time() - start_time
    load_mem = get_memory_usage() - start_mem
    
    print(f"📥 Load Time: {load_time:.4f}s")
    
    # Measure Clean Time (Dedupe + Fill Nulls + Standardization)
    start_time = time.time()
    
    # Simulate standard cleaning pipeline
    cleaned_df = (
        df.unique()  # Dedupe
        .fill_null(0)  # Simple fill
        .with_columns(pl.col("text_short").str.to_uppercase())  # String op
    )
    
    clean_time = time.time() - start_time
    clean_mem = get_memory_usage() - start_mem
    
    print(f"🧹 Clean Time: {clean_time:.4f}s")
    
    # Throughput
    rows = len(df)
    rows_per_sec = rows / clean_time
    print(f"⚡ Throughput: {rows_per_sec:,.0f} rows/sec")
    
    # Total Stats
    print("-" * 30)
    print(f"Total Time: {load_time + clean_time:.4f}s")
    print(f"Peak Memory Delta: {max(load_mem, clean_mem):.2f} MB")

def main():
    """Run benchmarks on all generated datasets."""
    datasets = [
        "benchmarks/data/dataset_small.csv",
        "benchmarks/data/dataset_medium.csv",
        # "benchmarks/data/dataset_large.csv"
    ]
    
    print("Starting Stress Tests...")
    print(f"CPU Cores: {os.cpu_count()}")
    print(f"Polars Version: {pl.__version__}")
    
    for ds in datasets:
        run_benchmark(ds)

if __name__ == "__main__":
    main()
