"""
TurboFrame Test Suite
Run: python -m pytest tests/test_turboframe.py -v
Or:  python tests/test_turboframe.py
"""

import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import pyarrow as pa
import numpy as np

from turboframe import TurboFrame


def test_create_from_dict():
    tf = TurboFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
    assert tf.shape == (3, 2)
    assert tf.columns == ["a", "b"]
    print("PASS: create from dict")


def test_create_from_list_of_dicts():
    data = [{"name": "Alice", "score": 95}, {"name": "Bob", "score": 87}]
    tf = TurboFrame(data)
    assert tf.shape == (2, 2)
    print("PASS: create from list of dicts")


def test_create_from_pandas():
    try:
        import pandas as pd
        pdf = pd.DataFrame({"x": [10, 20, 30], "y": ["a", "b", "c"]})
        tf = TurboFrame(pdf)
        assert tf.shape == (3, 2)
        print("PASS: create from pandas")
    except ImportError:
        print("SKIP: pandas not installed")


def test_create_from_arrow():
    table = pa.table({"col1": [1, 2, 3], "col2": ["x", "y", "z"]})
    tf = TurboFrame(table)
    assert len(tf) == 3
    print("PASS: create from arrow table")


def test_select_columns():
    tf = TurboFrame({"a": [1, 2], "b": [3, 4], "c": [5, 6]})
    result = tf.select(["a", "c"])
    assert result.columns == ["a", "c"]
    print("PASS: select columns")


def test_drop_columns():
    tf = TurboFrame({"a": [1, 2], "b": [3, 4], "c": [5, 6]})
    result = tf.drop("b")
    assert "b" not in result.columns
    print("PASS: drop columns")


def test_rename_columns():
    tf = TurboFrame({"a": [1, 2], "b": [3, 4]})
    result = tf.rename({"a": "alpha", "b": "beta"})
    assert result.columns == ["alpha", "beta"]
    print("PASS: rename columns")


def test_filter_numeric():
    tf = TurboFrame({"val": [10, 20, 30, 40, 50]})
    result = tf.filter("val > 25")
    assert len(result) == 3
    print("PASS: filter numeric")


def test_filter_string():
    tf = TurboFrame({"region": ["East", "West", "East"], "sales": [1, 2, 3]})
    result = tf.filter("region == 'East'")
    assert len(result) == 2
    print("PASS: filter string")


def test_sort():
    tf = TurboFrame({"val": [30, 10, 20]})
    result = tf.sort("val")
    assert result.to_dict()["val"] == [10, 20, 30]
    print("PASS: sort")


def test_groupby_sum():
    tf = TurboFrame({
        "region": ["E", "W", "E", "W", "E"],
        "sales":  [100, 200, 150, 300, 250],
    })
    result = tf.groupby("region").agg({"sales": "sum"})
    d = result.to_dict()
    for i, region in enumerate(d["region"]):
        if region == "E":
            assert d["sales_sum"][i] == 500
        elif region == "W":
            assert d["sales_sum"][i] == 500
    print("PASS: groupby sum")


def test_groupby_multiple_aggs():
    tf = TurboFrame({
        "region": ["E", "W", "E", "W"],
        "sales":  [100, 200, 300, 400],
        "qty":    [10, 20, 30, 40],
    })
    result = tf.groupby("region").agg({"sales": "sum", "qty": "mean"})
    assert "sales_sum" in result.columns
    assert "qty_mean" in result.columns
    print("PASS: groupby multiple aggs")


def test_join():
    left = TurboFrame({"id": [1, 2, 3], "name": ["a", "b", "c"]})
    right = TurboFrame({"id": [2, 3, 4], "score": [90, 80, 70]})
    result = left.join(right, on="id", how="inner")
    assert len(result) == 2
    print("PASS: join")


def test_unique():
    tf = TurboFrame({"region": ["E", "W", "E", "N", "W"]})
    assert set(tf.unique("region")) == {"E", "W", "N"}
    print("PASS: unique")


def test_value_counts():
    tf = TurboFrame({"region": ["E", "W", "E", "E", "W"]})
    vc = tf.value_counts("region")
    d = vc.to_dict()
    for i, r in enumerate(d["region"]):
        if r == "E":
            assert d["count"][i] == 3
    print("PASS: value_counts")


def test_to_pandas():
    tf = TurboFrame({"a": [1, 2, 3]})
    pdf = tf.to_pandas()
    assert len(pdf) == 3
    print("PASS: to_pandas")


def test_describe():
    tf = TurboFrame({"val": [10, 20, 30, 40, 50]})
    stats = tf.describe()
    assert stats["val"]["mean"] == 30.0
    print("PASS: describe")


def test_full_pipeline():
    tf = TurboFrame({
        "region":   ["East", "West", "East", "West", "East", "North"],
        "amount":   [100, 200, 150, 300, 250, 400],
        "quantity": [10, 20, 15, 30, 25, 40],
    })
    result = (
        tf.filter("amount > 100")
          .select(["region", "amount", "quantity"])
          .groupby("region")
          .agg({"amount": "sum", "quantity": "mean"})
          .sort("amount_sum", ascending=False)
    )
    assert len(result) > 0
    assert "amount_sum" in result.columns
    print("PASS: full pipeline")


def test_large_parallel_groupby():
    np.random.seed(42)
    n = 500_000
    regions = np.random.choice(["East", "West", "North", "South"], size=n)
    amounts = np.random.uniform(1, 1000, size=n)

    tf = TurboFrame({"region": regions.tolist(), "amount": amounts.tolist()})

    import time
    t0 = time.perf_counter()
    result = tf.groupby("region").agg({"amount": "sum"})
    elapsed = time.perf_counter() - t0

    assert len(result) == 4
    print(f"PASS: large parallel groupby ({n:,} rows in {elapsed:.3f}s)")


if __name__ == "__main__":
    print("=" * 60)
    print("TurboFrame Test Suite")
    print("=" * 60)

    tests = [
        test_create_from_dict, test_create_from_list_of_dicts,
        test_create_from_pandas, test_create_from_arrow,
        test_select_columns, test_drop_columns, test_rename_columns,
        test_filter_numeric, test_filter_string, test_sort,
        test_groupby_sum, test_groupby_multiple_aggs,
        test_join, test_unique, test_value_counts,
        test_to_pandas, test_describe,
        test_full_pipeline, test_large_parallel_groupby,
    ]

    passed = failed = 0
    for test in tests:
        try:
            test()
            passed += 1
        except Exception as e:
            print(f"FAIL: {test.__name__} - {e}")
            failed += 1

    print("=" * 60)
    print(f"Results: {passed} passed, {failed} failed, {len(tests)} total")
    print("=" * 60)
