# /// script
# requires-python = ">=3.12,<3.13"
# dependencies = ["pyarrow==25.0.1"]
# ///
import hashlib
import json
import sys
from decimal import Decimal
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

root = Path(__file__).resolve().parents[3]
target = root / 'explorer/static/arrow'
target.mkdir(parents=True, exist_ok=True)
arrays = {
    'id': pa.array([9007199254740993, 2, 3, 4, 5, 6], type=pa.int64()),
    'city': pa.array(['İzmir', None, '', '東京', 'Oslo', 'İzmir'], type=pa.string()),
    'readings': pa.array([[4, None, 7], None, [], [9], [2, 3], []], type=pa.list_(pa.int32())),
    'category': pa.DictionaryArray.from_arrays(pa.array([0, 1, None, 0, 1, 0], type=pa.int8()), pa.array(['coast', 'inland'])),
    'amount': pa.array([Decimal('12345678901234567890.12'), None, Decimal('0.00'), Decimal('-2.50'), Decimal('1.25'), Decimal('7.00')], type=pa.decimal128(24, 2)),
    'moment': pa.array([1700000000123456, None, 1700000000123458, 1700000000123459, 1700000000123460, 1700000000123461], type=pa.timestamp('us', tz='UTC')),
}
table = pa.table(arrays)


def describe(array):
    result = {
        'type': str(array.type), 'length': len(array), 'offset': array.offset,
        'nullCount': array.null_count,
        'buffers': [None if buffer is None else buffer.to_pybytes().hex() for buffer in array.buffers()[:array.type.num_buffers]],
        'values': [None if value is None else str(value) for value in array.to_pylist()],
    }
    if pa.types.is_list(array.type):
        result['child'] = describe(array.values)
    if pa.types.is_dictionary(array.type):
        result['dictionary'] = describe(array.dictionary)
    return result


artifacts = {}
for name, factory in [('sample.arrow', pa.ipc.new_file), ('sample.arrows', pa.ipc.new_stream)]:
    sink = pa.BufferOutputStream()
    with factory(sink, table.schema) as writer:
        writer.write_table(table, max_chunksize=3)
    payload = sink.getvalue().to_pybytes()
    restored = (pa.ipc.open_file if name.endswith('.arrow') else pa.ipc.open_stream)(pa.BufferReader(payload)).read_all()
    assert restored.equals(table)
    artifacts[name] = payload
sink = pa.BufferOutputStream()
pq.write_table(table, sink, row_group_size=3, compression='snappy', write_page_checksum=True)
payload = sink.getvalue().to_pybytes()
assert pq.read_table(pa.BufferReader(payload)).equals(table)
artifacts['sample.parquet'] = payload

fixture = {
    'writer': f'PyArrow {pa.__version__}', 'rows': table.num_rows, 'batchRows': [3, 3],
    'columns': {name: describe(array) for name, array in arrays.items()},
    'slice': {name: describe(array.slice(1, 3)) for name, array in arrays.items()},
    'files': {name: {'size': len(payload), 'sha256': hashlib.sha256(payload).hexdigest()} for name, payload in artifacts.items()},
}
assert arrays['city'].buffers()[1].to_pybytes().hex() == '000000000600000006000000060000000c0000001000000016000000'
assert arrays['readings'].offsets.to_pylist() == [0, 3, 3, 3, 4, 6, 6]
assert arrays['readings'].values.to_pylist() == [4, None, 7, 9, 2, 3]
assert arrays['city'].slice(1, 3).offset == 1
assert arrays['city'].slice(1, 3).buffers()[2].address == arrays['city'].buffers()[2].address
assert arrays['id'][0].as_py() == 9007199254740993

verify = '--verify' in sys.argv
for name, payload in artifacts.items():
    path = target / name
    if verify:
        assert path.read_bytes() == payload, name
    else:
        path.write_bytes(payload)
for path in [target / 'buffers.json', root / 'explorer/src/lib/arrow/generated.json']:
    if verify:
        assert json.loads(path.read_text()) == fixture, path
    else:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(fixture, indent=2, ensure_ascii=False) + '\n')
print('Verified Arrow file/stream and Parquet parity, exact integers, UTF-8/list offsets and shared-buffer slicing.')