#!/usr/bin/env python3
"""Test_PyArrow.py - on-device validation for the pyarrow 25.0.1 wheel.

Exercises arrays, tables, IPC, CSV, JSON, Feather, compute, and types.
Exit code 0 = all tests passed.

Generated by RIMI
"""
import os
import sys
import tempfile

RESULTS = []


def test(name, fn):
    try:
        fn()
        RESULTS.append(("PASS", name))
    except NotImplementedError:
        RESULTS.append(("SKIP", name))
    except Exception as e:  # noqa: BLE001
        RESULTS.append(("FAIL", name, str(e)))


def section(title):
    print("\n===== %s =====" % title)


def check(cond, msg="assertion failed"):
    if not cond:
        raise AssertionError(msg)


WORKDIR = None


def workdir():
    global WORKDIR
    if WORKDIR is None:
        import os as _os
        candidates = [_os.environ.get("TMPDIR") or "", tempfile.gettempdir(),
                      "/storage/emulated/0/Download", _os.getcwd()]
        for base in candidates:
            if not base:
                continue
            try:
                d = os.path.join(base, "test_pyarrow_tmp")
                os.makedirs(d, exist_ok=True)
                with open(os.path.join(d, "_probe"), "w") as fh:
                    fh.write("ok")
                WORKDIR = d
                break
            except OSError:
                continue
        if WORKDIR is None:
            WORKDIR = "."
    return WORKDIR


# ---------------------------------------------------------------------------
# 1. import / version
# ---------------------------------------------------------------------------
def test_import_pyarrow():
    import pyarrow as pa
    print("    pyarrow", pa.__version__)
    check(pa.__version__.startswith("25."), "unexpected version: %s" % pa.__version__)


def test_import_numpy_dep():
    import numpy as np
    print("    numpy", np.__version__)
    check(hasattr(np, "ndarray"), "numpy not functional")


def test_import_submodules():
    import pyarrow.compute as pc
    import pyarrow.csv
    import pyarrow.json
    import pyarrow.feather
    import pyarrow.fs
    import pyarrow.ipc
    check(hasattr(pc, "cast"), "compute module incomplete")
    print("    compute, csv, json, feather, fs, ipc -- all imported")


# ---------------------------------------------------------------------------
# 2. arrays
# ---------------------------------------------------------------------------
def test_array_creation():
    import pyarrow as pa
    a = pa.array([1, 2, 3, 4, 5])
    check(a.type == pa.int64(), "type %r" % a.type)
    check(len(a) == 5, "len %d" % len(a))
    check(a.to_pylist() == [1, 2, 3, 4, 5], "values mismatch")
    print("    int64 array:", a)


def test_array_float():
    import pyarrow as pa
    a = pa.array([1.0, 2.5, 3.7], type=pa.float64())
    check(a.type == pa.float64(), "type %r" % a.type)
    check(len(a) == 3, "len %d" % len(a))
    print("    float64 array:", a)


def test_array_string():
    import pyarrow as pa
    a = pa.array(["hello", "world", "pyarrow"])
    check(a.type == pa.string(), "type %r" % a.type)
    check(a.to_pylist() == ["hello", "world", "pyarrow"])
    print("    string array:", a)


def test_array_null():
    import pyarrow as pa
    a = pa.array([1, None, 3], type=pa.int64())
    check(a.null_count == 1, "null_count %d" % a.null_count)
    check(a.to_pylist() == [1, None, 3])
    print("    null handling:", a)


# ---------------------------------------------------------------------------
# 3. types
# ---------------------------------------------------------------------------
def test_types():
    import pyarrow as pa
    t_int = pa.int64()
    t_float = pa.float64()
    t_str = pa.string()
    t_bool = pa.bool_()
    check(str(t_int) == "int64", "int64 str: %r" % str(t_int))
    check(str(t_float) in ("float64", "double"), "float64 str: %r" % str(t_float))
    check(str(t_str) == "string", "string str: %r" % str(t_str))
    check(str(t_bool) == "bool", "bool str: %r" % str(t_bool))
    check(isinstance(t_int, pa.DataType), "not DataType")
    print("    types: int64, float64, string, bool -- all recognized")


# ---------------------------------------------------------------------------
# 4. tables
# ---------------------------------------------------------------------------
def test_table_creation():
    import pyarrow as pa
    t = pa.table({
        "id": [1, 2, 3],
        "name": ["alice", "bob", "charlie"],
        "score": [95.5, 87.0, 92.3],
    })
    check(t.num_rows == 3, "rows %d" % t.num_rows)
    check(t.num_columns == 3, "cols %d" % t.num_columns)
    check(t.column_names == ["id", "name", "score"])
    check(t.schema.field("id").type == pa.int64())
    check(t.schema.field("name").type == pa.string())
    check(t.schema.field("score").type == pa.float64())
    print("    table: %d rows x %d cols" % (t.num_rows, t.num_columns))


def test_table_ops():
    import pyarrow as pa
    t = pa.table({"x": [10, 20, 30], "y": [1.0, 2.0, 3.0]})
    check(t.column("x").to_pylist() == [10, 20, 30])
    check(t.to_pandas().shape == (3, 2))
    print("    table column access + to_pandas OK")


# ---------------------------------------------------------------------------
# 5. IPC round-trip
# ---------------------------------------------------------------------------
def test_ipc_roundtrip():
    import pyarrow as pa
    import pyarrow.ipc as ipc
    t = pa.table({
        "a": [1, 2, 3, 4, 5],
        "b": ["x", "y", "z", "w", "v"],
        "c": [1.1, 2.2, 3.3, 4.4, 5.5],
    })
    path = os.path.join(workdir(), "test_ipc.arrow")
    sink = ipc.new_file(path, t.schema)
    sink.write_table(t)
    sink.close()
    reader = ipc.open_file(path)
    t2 = reader.read_all()
    check(t.equals(t2), "IPC roundtrip mismatch")
    os.unlink(path)
    print("    IPC file: write -> read -> equals")


def test_ipc_stream():
    import pyarrow as pa
    import pyarrow.ipc as ipc
    t = pa.table({"val": [100, 200, 300]})
    path = os.path.join(workdir(), "test_ipc_stream.arrow")
    sink = ipc.new_stream(path, t.schema)
    sink.write_table(t)
    sink.close()
    reader = ipc.open_stream(path)
    t2 = reader.read_all()
    check(t.equals(t2), "IPC stream roundtrip mismatch")
    os.unlink(path)
    print("    IPC stream: write -> read -> equals")


# ---------------------------------------------------------------------------
# 6. CSV
# ---------------------------------------------------------------------------
def test_csv_write_read():
    import pyarrow as pa
    import pyarrow.csv as pcsv
    t = pa.table({
        "id": [1, 2, 3],
        "name": ["alice", "bob", "charlie"],
        "val": [10.5, 20.3, 30.1],
    })
    path = os.path.join(workdir(), "test_csv.csv")
    pcsv.write_csv(t, path)
    t2 = pcsv.read_csv(path)
    check(t2.num_rows == 3, "rows %d" % t2.num_rows)
    check(t2.num_columns == 3, "cols %d" % t2.num_columns)
    os.unlink(path)
    print("    CSV write -> read: %d rows x %d cols" % (t2.num_rows, t2.num_columns))


def test_csv_options():
    import pyarrow as pa
    import pyarrow.csv as pcsv
    path = os.path.join(workdir(), "test_csv_opts.csv")
    with open(path, "w") as fh:
        fh.write("1,2,3\n4,5,6\n")
    read_opts = pcsv.ReadOptions(column_names=["x", "y", "z"])
    convert_opts = pcsv.ConvertOptions(column_types={"x": pa.int64(), "y": pa.int64(), "z": pa.int64()})
    t = pcsv.read_csv(path, read_options=read_opts, convert_options=convert_opts)
    check(t.column("x").to_pylist() == [1, 4])
    check(t.schema.field("x").type == pa.int64())
    os.unlink(path)
    print("    CSV with custom options: column_names + column_types")


# ---------------------------------------------------------------------------
# 7. JSON
# ---------------------------------------------------------------------------
def test_json_read():
    import pyarrow as pa
    import pyarrow.json as pjson
    path = os.path.join(workdir(), "test_json.json")
    with open(path, "w") as fh:
        fh.write('{"a": 1, "b": "hello"}\n')
        fh.write('{"a": 2, "b": "world"}\n')
    t = pjson.read_json(path)
    check(t.num_rows == 2, "rows %d" % t.num_rows)
    check("a" in t.column_names, "column 'a' missing")
    check("b" in t.column_names, "column 'b' missing")
    os.unlink(path)
    print("    JSON read: %d rows, columns=%r" % (t.num_rows, t.column_names))


# ---------------------------------------------------------------------------
# 8. Feather round-trip
# ---------------------------------------------------------------------------
def test_feather_roundtrip():
    import pyarrow as pa
    import pyarrow.feather as pf
    t = pa.table({
        "id": [1, 2, 3, 4, 5],
        "name": ["alice", "bob", "charlie", "diana", "eve"],
        "score": [95.5, 87.0, 92.3, 88.8, 99.1],
    })
    path = os.path.join(workdir(), "test_feather.feather")
    pf.write_feather(t, path)
    t2 = pf.read_table(path)
    check(t.equals(t2), "Feather roundtrip mismatch")
    check(t2.num_rows == 5, "rows %d" % t2.num_rows)
    os.unlink(path)
    print("    Feather: write -> read -> equals (%d rows)" % t2.num_rows)


# ---------------------------------------------------------------------------
# 9. compute
# ---------------------------------------------------------------------------
def test_compute_basic():
    import pyarrow as pa
    import pyarrow.compute as pc
    a = pa.array([1, 2, 3, 4, 5])
    result = pc.sum(a)
    check(result.as_py() == 15, "sum %r" % result)
    print("    compute.sum:", result.as_py())


def test_compute_cast():
    import pyarrow as pa
    import pyarrow.compute as pc
    a = pa.array([1, 2, 3], type=pa.int64())
    b = pc.cast(a, pa.float64())
    check(b.type == pa.float64(), "type %r" % b.type)
    check(b.to_pylist() == [1.0, 2.0, 3.0])
    print("    compute.cast int64 -> float64:", b)


def test_compute_filter():
    import pyarrow as pa
    import pyarrow.compute as pc
    a = pa.array([10, 20, 30, 40, 50])
    mask = pc.greater(a, 25)
    filtered = pc.filter(a, mask)
    check(filtered.to_pylist() == [30, 40, 50])
    print("    compute.filter > 25:", filtered)


def test_compute_arithmetic():
    import pyarrow as pa
    import pyarrow.compute as pc
    a = pa.array([10, 20, 30])
    b = pa.array([1, 2, 3])
    add_result = pc.add(a, b)
    mul_result = pc.multiply(a, b)
    check(add_result.to_pylist() == [11, 22, 33])
    check(mul_result.to_pylist() == [10, 40, 90])
    print("    compute add/multiply:", add_result, mul_result)


# ---------------------------------------------------------------------------
# 10. fs (filesystem)
# ---------------------------------------------------------------------------
def test_fs_local():
    import pyarrow.fs as pfs
    local = pfs.LocalFileSystem()
    path = os.path.join(workdir(), "test_fs.txt")
    with open(path, "w") as fh:
        fh.write("filesystem test")
    meta = local.get_file_info(path)
    check(meta.type == pfs.FileType.File, "not a file")
    check(meta.size > 0, "size %d" % meta.size)
    os.unlink(path)
    print("    LocalFileSystem: get_file_info OK (size=%d)" % meta.size)


# ---------------------------------------------------------------------------
# def main
# ---------------------------------------------------------------------------
def main():
    section("pyarrow 25.0.1 - import / basics")
    test("import pyarrow (25.x)", test_import_pyarrow)
    test("import numpy (dependency)", test_import_numpy_dep)
    test("import submodules (compute, csv, json, feather, fs, ipc)", test_import_submodules)

    section("arrays")
    test("array int64", test_array_creation)
    test("array float64", test_array_float)
    test("array string", test_array_string)
    test("array null handling", test_array_null)

    section("types")
    test("types (int64, string, float64, bool)", test_types)

    section("tables")
    test("table creation + schema", test_table_creation)
    test("table column access + to_pandas", test_table_ops)

    section("IPC")
    test("IPC file round-trip", test_ipc_roundtrip)
    test("IPC stream round-trip", test_ipc_stream)

    section("CSV")
    test("CSV write / read", test_csv_write_read)
    test("CSV custom options", test_csv_options)

    section("JSON")
    test("JSON read", test_json_read)

    section("Feather")
    test("Feather round-trip", test_feather_roundtrip)

    section("compute")
    test("compute.sum", test_compute_basic)
    test("compute.cast", test_compute_cast)
    test("compute.filter", test_compute_filter)
    test("compute arithmetic", test_compute_arithmetic)

    section("filesystem")
    test("LocalFileSystem get_file_info", test_fs_local)

    section("RESULT")
    n_ok = n_fail = n_skip = 0
    for r in RESULTS:
        status = r[0]
        if status == "PASS":
            n_ok += 1
            print("  OK    %s" % r[1])
        elif status == "SKIP":
            n_skip += 1
            print("  SKIP  %s" % r[1])
        else:
            n_fail += 1
            print("  FAIL  %s: %s" % (r[1], r[2]))
    print("RESULT: %d ok, %d failed, %d skipped" % (n_ok, n_fail, n_skip))
    sys.exit(1 if n_fail else 0)


if __name__ == "__main__":
    main()