from __future__ import annotations

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

import snapparquet as sp


def test_schema_and_intersect(tmp_path):
    path = tmp_path / "t.parquet"
    df = pd.DataFrame(
        {
            "load_date": [20250101, 20250102, 20250315],
            "a": [1, 2, 3],
            "b": [10, 20, 30],
        }
    )
    table = pa.Table.from_pandas(df)
    pq.write_table(table, path)

    names = sp.schema_column_names(str(path))
    assert names == {"load_date", "a", "b"}
    assert sp.column_is_integer_type(str(path), "load_date")
    sub = sp.intersect_columns(str(path), ["b", "load_date", "missing"])
    assert sub == ["b", "load_date"]


def test_read_parquet_prune_and_filter(tmp_path):
    path = tmp_path / "t.parquet"
    df = pd.DataFrame(
        {
            "load_date": [20250101, 20250201, 20250301],
            "x": [1, 2, 3],
            "noise": ["a", "b", "c"],
        }
    )
    pq.write_table(pa.Table.from_pandas(df), path)

    out = sp.read_parquet(
        str(path),
        columns=["load_date", "x"],
        filter_column="load_date",
        range_start=20250201,
        range_end=20250301,
    )
    assert list(out.columns) == ["load_date", "x"]
    assert len(out) == 2
    assert out["x"].tolist() == [2, 3]


def test_read_tabular_many(tmp_path):
    p1 = tmp_path / "d1.parquet"
    p2 = tmp_path / "d2.parquet"
    pq.write_table(
        pa.Table.from_pandas(pd.DataFrame({"load_date": [20250101], "v": [1]})),
        p1,
    )
    pq.write_table(
        pa.Table.from_pandas(pd.DataFrame({"load_date": [20250102], "v": [2]})),
        p2,
    )
    out = sp.read_tabular_many(
        [str(p1), str(p2)],
        columns=["load_date", "v"],
    )
    assert len(out) == 2
    assert set(out["v"].tolist()) == {1, 2}


def test_read_tabular_many_unified_pushdown_and_sequential(tmp_path):
    p1 = tmp_path / "d1.parquet"
    p2 = tmp_path / "d2.parquet"
    pq.write_table(
        pa.Table.from_pandas(
            pd.DataFrame({"load_date": [20250101, 20250201], "v": [1, 2]})
        ),
        p1,
    )
    pq.write_table(
        pa.Table.from_pandas(
            pd.DataFrame({"load_date": [20250301, 20250401], "v": [3, 4]})
        ),
        p2,
    )
    out = sp.read_tabular_many(
        [str(p1), str(p2)],
        columns=["load_date", "v"],
        filter_column="load_date",
        range_start=20250201,
        range_end=20250301,
    )
    assert len(out) == 2
    assert set(out["v"].tolist()) == {2, 3}

    out_seq = sp.read_tabular_many(
        [str(p1), str(p2)],
        columns=["load_date", "v"],
        unified_parquet=False,
        parallel_workers=1,
    )
    assert len(out_seq) == 4


def test_schema_matches_any_group():
    names = frozenset({"pdp_atg", "recommended_atg", "load_date"})
    groups = (["pivot_atg_level", "shown_atg_level"], ["pdp_atg", "recommended_atg"])
    assert sp.schema_matches_any_group(names, groups)


def test_build_filter_and_read_parquet_filters(tmp_path):
    path = tmp_path / "f.parquet"
    df0 = pd.DataFrame(
        {
            "status": ["active", "trial", "done"],
            "revenue": [10, 0, 5],
            "id": [1, 2, 3],
        }
    )
    pq.write_table(pa.Table.from_pandas(df0), path)

    expr = sp.build_filter(
        [("status", "in", ["active", "trial"]), ("revenue", ">", 0)]
    )
    assert expr is not None

    out = sp.read_parquet(
        str(path),
        columns=["status", "revenue"],
        filters=[("status", "in", ["active", "trial"]), ("revenue", ">", 0)],
    )
    assert len(out) == 1
    assert out["status"].iloc[0] == "active"
    assert int(out["revenue"].iloc[0]) == 10


def test_read_parquet_filters_plus_legacy_range(tmp_path):
    path = tmp_path / "fr.parquet"
    df0 = pd.DataFrame(
        {
            "load_date": [20250101, 20250201, 20250301],
            "region": ["A", "B", "A"],
        }
    )
    pq.write_table(pa.Table.from_pandas(df0), path)
    out = sp.read_parquet(
        str(path),
        columns=["load_date", "region"],
        filter_column="load_date",
        range_start=20250201,
        range_end=20250301,
        filters=[("region", "==", "A")],
    )
    assert len(out) == 1
    assert int(out["load_date"].iloc[0]) == 20250301


def test_read_parquet_filters_plus_legacy_range_when_range_not_pushdown(tmp_path):
    """Legacy range on non-integer Parquet column: filters push down; range must still apply."""
    path = tmp_path / "fr_str.parquet"
    df0 = pd.DataFrame(
        {
            "load_date": ["20250101", "20250201", "20250301"],
            "region": ["A", "B", "A"],
        }
    )
    pq.write_table(pa.Table.from_pandas(df0), path)
    out = sp.read_parquet(
        str(path),
        columns=["load_date", "region"],
        filter_column="load_date",
        range_start=20250201,
        range_end=20250301,
        filters=[("region", "==", "A")],
    )
    assert len(out) == 1
    assert out["region"].iloc[0] == "A"
    assert int(pd.to_numeric(out["load_date"].iloc[0])) == 20250301


def test_apply_filters_pandas_missing_column_raises():
    df = pd.DataFrame({"a": [1]})
    try:
        sp.apply_filters_pandas(df, [("missing", "==", 1)])
    except ValueError as e:
        assert "unknown column" in str(e).lower()
    else:
        raise AssertionError("expected ValueError")


def test_dictionary_encoded_to_categorical(tmp_path):
    path = tmp_path / "dict.parquet"
    idx = pa.array([0, 1, 0, 2], type=pa.int8())
    dictionary = pa.array(["x", "y", "z"])
    dict_arr = pa.DictionaryArray.from_arrays(idx, dictionary)
    tbl = pa.table({"tag": dict_arr, "v": [1, 2, 3, 4]})
    pq.write_table(tbl, path)
    out = sp.read_parquet(
        str(path),
        columns=["tag", "v"],
        backend="numpy",
        dictionary_as_categorical=True,
    )
    assert str(out["tag"].dtype) == "category"
    assert out["tag"].tolist() == ["x", "y", "x", "z"]
