"""Round-trip every exported payload through a real Arrow consumer.

The node builds its Arrow arrays by hand against the C data interface rather
than linking Arrow C++, so this is what shows that what dora imports on the
other side of `send_arrow_output` is what went in. It needs neither a camera nor
the ZED SDK.
"""

import ctypes
import pathlib
import sys

import pyarrow as pa

LIBRARY = pathlib.Path(__file__).resolve().parent.parent / "build" / "tests" / "libzed_arrow.so"

# ArrowArray is ten machine words and ArrowSchema nine, and both want the
# 8-byte alignment an array of uint64 gives.
WORDS = 16
# Index of the `release` field in ArrowArray: a consumer that took ownership
# zeroes it, which is how this test sees the move happen.
RELEASE = 8

IMAGE_TYPE = pa.struct([
    ("width", pa.uint32()),
    ("height", pa.uint32()),
    ("encoding", pa.utf8()),
    ("data", pa.large_binary()),
])


def load():
    """Load the test library and declare the argument types of its exports."""
    library = ctypes.CDLL(str(LIBRARY))

    library.zed_export_image.argtypes = [
        ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32,
        ctypes.c_char_p, ctypes.c_char_p, ctypes.c_size_t,
    ]
    library.zed_export_float32.argtypes = [
        ctypes.c_void_p, ctypes.c_void_p, ctypes.POINTER(ctypes.c_float), ctypes.c_size_t,
    ]
    library.zed_export_string.argtypes = [
        ctypes.c_void_p, ctypes.c_void_p, ctypes.c_char_p, ctypes.c_size_t,
    ]
    return library


def structs():
    """Allocate the ArrowArray and ArrowSchema the producer fills in."""
    return (ctypes.c_uint64 * WORDS)(), (ctypes.c_uint64 * WORDS)()


def imported(array, schema):
    """Import the pair as pyarrow does, taking ownership of both."""
    result = pa.Array._import_from_c(ctypes.addressof(array), ctypes.addressof(schema))
    assert array[RELEASE] == 0, "the consumer should have taken ownership of the array"
    return result


def test_image(library):
    """An image row carries its own dimensions, encoding and bytes."""
    array, schema = structs()
    pixels = bytes(range(12))

    library.zed_export_image(
        ctypes.byref(array), ctypes.byref(schema), 2, 2, b"bgr8", pixels, len(pixels)
    )
    result = imported(array, schema)

    assert result.type == IMAGE_TYPE, result.type
    assert result.to_pylist() == [
        {"width": 2, "height": 2, "encoding": "bgr8", "data": pixels}
    ]


def test_depth(library):
    """A depth map is an image whose bytes are float32 millimetres."""
    array, schema = structs()
    millimetres = [1500.0, 0.0, 250.5, 4000.0]
    payload = pa.array(millimetres, type=pa.float32()).buffers()[1].to_pybytes()

    library.zed_export_image(
        ctypes.byref(array), ctypes.byref(schema), 2, 2, b"32FC1", payload, len(payload)
    )
    row = imported(array, schema).to_pylist()[0]

    assert row["encoding"] == "32FC1"
    assert pa.Array.from_buffers(
        pa.float32(), 4, [None, pa.py_buffer(row["data"])]
    ).to_pylist() == millimetres


def test_float32(library):
    """Intrinsics and the point cloud go out as plain float32 values."""
    array, schema = structs()
    values = (ctypes.c_float * 4)(700.5, 700.5, 640.0, 360.0)

    library.zed_export_float32(ctypes.byref(array), ctypes.byref(schema), values, 4)
    result = imported(array, schema)

    assert result.type == pa.float32()
    assert result.to_pylist() == [700.5, 700.5, 640.0, 360.0]


def test_string(library):
    """The camera state goes out as one UTF-8 string."""
    array, schema = structs()

    library.zed_export_string(ctypes.byref(array), ctypes.byref(schema), b"streaming", 9)
    result = imported(array, schema)

    assert result.type == pa.utf8()
    assert result.to_pylist() == ["streaming"]


def main():
    """Run every test, reporting the first failure."""
    library = load()

    for test in (test_image, test_depth, test_float32, test_string):
        test(library)
        print(f"ok  {test.__name__}")

    return 0


if __name__ == "__main__":
    sys.exit(main())