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_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)
