# /// script
# requires-python = ">=3.12,<3.13"
# dependencies = ["pylance==11.0.0", "pyarrow==25.0.1", "numpy==2.4.3", "protobuf==6.33.5"]
# ///
import argparse
import hashlib
import json
import shutil
import struct
import tempfile
import zipfile
from pathlib import Path

import lance
import numpy as np
import pyarrow as pa
from google.protobuf.any_pb2 import Any
from google.protobuf.json_format import MessageToDict
from lance.file import LanceFileReader, LanceFileWriter
from lance.fragment import DataFile, FragmentMetadata
from encodings_v2_1_pb2 import PageLayout, REPDEF_ALL_VALID_ITEM, REPDEF_NULLABLE_ITEM, COMPRESSION_ALGORITHM_ZSTD
from file2_pb2 import ColumnMetadata


COMMIT = 'ab6b5bbe46009ed78746b444df8db59a8bc5d842'
SELECTED = [0, 6, 7, 255, 256, 1000, 1024, 2047]


def digest(raw):
    return hashlib.sha256(raw).hexdigest()


def source_table(compression):
    schema = pa.schema([
        pa.field('id', pa.int64(), metadata={'lance-encoding:structural-encoding': 'fullzip'}),
        pa.field('label', pa.string(), metadata={'lance-encoding:compression': compression}),
    ])
    labels = [None if index % 7 == 0 else f'row-{index}-' + 'abcdef' * 20 for index in range(2048)]
    return pa.Table.from_arrays([pa.array(np.arange(2048, dtype=np.int64)), pa.array(labels)], schema=schema)


def inspect(directory, compression):
    dataset = lance.dataset(directory)
    original = source_table(compression)
    assert dataset.to_table().equals(original)
    fragments = [fragment.metadata.to_json() for fragment in dataset.get_fragments()]
    assert len(fragments) == 1 and fragments[0]['id'] == 0
    assert fragments[0]['row_id_meta'] is None and fragments[0]['deletion_file'] is None
    assert len(fragments[0]['files']) == 1 and fragments[0]['physical_rows'] == 2048
    reference = fragments[0]['files'][0]
    assert reference['path'] == 'pages.lance' and reference['fields'] == reference['column_indices'] == [0, 1]
    file = directory / 'data/pages.lance'
    raw = file.read_bytes()
    sdk = LanceFileReader(str(file)).metadata()
    metadata_start, column_table, global_table, globals_count, columns_count, major, minor, magic = struct.unpack('<QQQIIHH4s', raw[-40:])
    assert magic == b'LANC' and (major, minor) == (2, 1)
    assert columns_count == len(sdk.columns) == 2 and globals_count == len(sdk.global_buffers)
    columns = []
    for index in range(columns_count):
        position, size = struct.unpack_from('<QQ', raw, column_table + index * 16)
        assert metadata_start <= position < position + size <= column_table
        metadata = ColumnMetadata.FromString(raw[position:position + size])
        assert len(metadata.pages) == len(sdk.columns[index].pages) > 1
        wrapper = Any.FromString(metadata.encoding.direct.encoding)
        assert wrapper.type_url == '/lance.encodings.ColumnEncoding' and wrapper.value == b'\x0a\x00'
        pages = []
        start = 0
        for number, page in enumerate(metadata.pages):
            assert page.priority == start and page.length > 0
            start += page.length
            assert page.encoding.WhichOneof('location') == 'direct'
            envelope = Any.FromString(page.encoding.direct.encoding)
            assert envelope.type_url == '/lance.encodings21.PageLayout'
            layout = PageLayout.FromString(envelope.value)
            buffers = [{'position': offset, 'size': length} for offset, length in zip(page.buffer_offsets, page.buffer_sizes, strict=True)]
            assert buffers == [{'position': buffer.position, 'size': buffer.size} for buffer in sdk.columns[index].pages[number].buffers]
            assert all(0 <= buffer['position'] <= buffer['position'] + buffer['size'] <= metadata_start for buffer in buffers)
            if index == 0:
                assert layout.WhichOneof('layout') == 'full_zip_layout'
                zipped = layout.full_zip_layout
                assert zipped.bits_rep == zipped.bits_def == 0 and zipped.bits_per_value == 64
                assert zipped.num_items == zipped.num_visible_items == page.length
                assert list(zipped.layers) == [REPDEF_ALL_VALID_ITEM]
                assert zipped.value_compression.WhichOneof('compression') == 'flat'
                assert zipped.value_compression.flat.bits_per_value == 64 and not zipped.value_compression.flat.HasField('data')
                assert len(buffers) == 1 and buffers[0]['size'] == page.length * 8
                values = np.frombuffer(raw, dtype='<i8', count=page.length, offset=buffers[0]['position'])
                np.testing.assert_array_equal(values, np.arange(page.priority, page.priority + page.length))
            else:
                assert layout.WhichOneof('layout') == 'mini_block_layout'
                block = layout.mini_block_layout
                assert list(block.layers) == [REPDEF_NULLABLE_ITEM] and block.num_items == page.length
                values = block.value_compression
                if compression == 'zstd':
                    assert values.WhichOneof('compression') == 'general'
                    assert values.general.compression.scheme == COMPRESSION_ALGORITHM_ZSTD
                    values = values.general.values
                assert values.WhichOneof('compression') == 'variable'
                assert values.variable.offsets.flat.bits_per_value == 32
            pages.append({'page': number, 'firstRow': page.priority, 'rows': page.length, 'buffers': buffers, 'layout': MessageToDict(layout, preserving_proto_field_name=True)})
        assert start == sdk.num_rows == original.num_rows
        columns.append({'name': original.schema[index].name, 'fieldId': fragments[0]['files'][0]['fields'][index], 'metadata': {'position': position, 'size': size}, 'pages': pages})
    queries = []
    for identity in SELECTED:
        result = dataset.to_table(filter=f'id = {identity}', with_row_address=True).to_pylist()
        assert result == [{**original.slice(identity, 1).to_pylist()[0], '_rowaddr': identity}]
        assert LanceFileReader(str(file)).read_range(identity, 1).to_table().to_pylist() == original.slice(identity, 1).to_pylist()
        queries.append({'id': identity, 'rowAddress': str(identity), 'query': f'id = {identity}', 'label': result[0]['label'], 'pages': [next(page['page'] for page in column['pages'] if page['firstRow'] <= identity < page['firstRow'] + page['rows']) for column in columns]})
    return {'compression': compression, 'version': str(dataset.version), 'fragment': '0', 'file': compression + '/data/pages.lance',
            'size': len(raw), 'sha256': digest(raw), 'rows': original.num_rows, 'nulls': original['label'].null_count,
            'manifest': compression + f'/_versions/{2 ** 64 - 1 - dataset.version}.manifest',
            'metadataStart': metadata_start, 'columnTable': column_table, 'globalTable': global_table,
            'columns': columns, 'queries': queries}


def verify(output, proof):
    assert proof['schemaVersion'] == 1 and proof['sourceCommit'] == COMMIT
    for path, item in proof['artifacts'].items():
        raw = (output / path).read_bytes()
        assert len(raw) == item['size'] and digest(raw) == item['sha256'], path
    with zipfile.ZipFile(output / 'originals.zip') as archive:
        assert set(archive.namelist()) == {path for path in proof['artifacts'] if path.startswith(('none/', 'zstd/'))}
        for path in archive.namelist():
            assert archive.read(path) == (output / path).read_bytes()
    for item in proof['layouts']:
        assert inspect(output / item['compression'], item['compression']) == item
    assert proof['layouts'][1]['size'] < proof['layouts'][0]['size']


parser = argparse.ArgumentParser()
parser.add_argument('--verify', action='store_true')
parser.add_argument('--output', type=Path, default=Path('explorer/static/lance/physical'))
args = parser.parse_args()
output = args.output.resolve()
pa.set_cpu_count(2)
assert lance.__version__ == '11.0.0'
if not args.verify and output.exists():
    raise RuntimeError('Inspect the existing physical-layout evidence before replacing it.')
with tempfile.TemporaryDirectory(prefix='columnar-lance-pages-') as work:
    root = Path(work)
    layouts = []
    for compression in ['none', 'zstd']:
        original = source_table(compression)
        directory = root / compression
        dataset = lance.write_dataset(original.slice(0, 0), directory, data_storage_version='2.1', enable_stable_row_ids=False)
        file = directory / 'data/pages.lance'
        file.parent.mkdir(exist_ok=True)
        with LanceFileWriter(str(file), schema=original.schema, version='2.1', data_cache_bytes=4096, max_page_bytes=4096) as writer:
            for batch in original.to_batches(max_chunksize=128):
                writer.write_batch(batch)
        fragment = FragmentMetadata(id=0, files=[DataFile.create(dataset, 'pages.lance')], physical_rows=original.num_rows)
        lance.LanceDataset.commit(dataset, lance.LanceOperation.Append([fragment]), read_version=dataset.version)
        layouts.append(inspect(directory, compression))
    if args.verify:
        proof = json.loads((output / 'proof.json').read_text())
        verify(output, proof)
        assert layouts == proof['layouts']
        assert (output / 'reproduce.py').read_bytes() == Path(__file__).read_bytes()
        for name in ['file2_pb2.py', 'encodings_v2_1_pb2.py', 'protos/file2.proto', 'protos/encodings_v2_1.proto', 'protos/LICENSE']:
            assert (output / name).read_bytes() == (Path(__file__).parent / name).read_bytes()
        print('Verified two original Lance layouts, every flat INT64 value from original bytes, 62 protobuf pages, complete nullable rows, sixteen native row selections and fresh byte-identical data files.')
    else:
        shutil.copytree(root, output)
        with zipfile.ZipFile(output / 'originals.zip', 'w', compression=zipfile.ZIP_DEFLATED) as archive:
            for path in sorted(output.rglob('*')):
                if path.is_file() and path.suffix != '.zip':
                    archive.write(path, path.relative_to(output))
        source = Path(__file__).parent
        for name in ['file2_pb2.py', 'encodings_v2_1_pb2.py']:
            shutil.copyfile(source / name, output / name)
        (output / 'protos').mkdir()
        for name in ['file2.proto', 'encodings_v2_1.proto', 'LICENSE']:
            shutil.copyfile(source / 'protos' / name, output / 'protos' / name)
        artifacts = {str(path.relative_to(output)): {'size': path.stat().st_size, 'sha256': digest(path.read_bytes())}
                     for path in sorted(output.rglob('*')) if path.is_file()}
        proof = {'schemaVersion': 1, 'sdk': 'pylance 11.0.0', 'format': '2.1', 'sourceCommit': COMMIT,
                 'oracle': 'Independent source arrays; NumPy little-endian INT64 decoding; pinned protobuf metadata',
                 'layouts': layouts, 'artifacts': artifacts}
        (output / 'proof.json').write_text(json.dumps(proof, indent=2) + '\n')
        shutil.copyfile(__file__, output / 'reproduce.py')
        (output / 'README.md').write_text(f'''# Multiple pages, nullable strings and compression\n\nTwo native Lance 2.1 tables contain the same 2,048 IDs and nullable strings. Each has\n15 ID pages and 16 string pages. The ID column deliberately uses all-valid FullZip\nflat INT64; string pages use MiniBlock with either no general compression or Zstd.\nThe browser decoder supports only that exact ID encoding. String values are native\nSDK results checked against independent source values, not browser-decoded strings.\n\nDownload this directory, extract originals.zip beside proof.json and reproduce.py,\nand retain both generated *_pb2.py modules and the protos directory. Run:\n\n```sh\nuv run reproduce.py --verify --output .\n```\n\nThe original protobuf schemas are Apache-2.0 licensed by the Lance authors, from\nhttps://github.com/lance-format/lance/tree/{COMMIT}/protos . Generated readers use\ngrpcio-tools 1.76.0 and protobuf 6.33.5. Original hashes and all native rows are checked;\na fresh writer must reproduce the two data files byte-for-byte. Manifest UUIDs and\ntimestamps may differ. Whole file sizes describe these examples, not an I/O benchmark.\n''')
        verify(output, proof)
        print(f'Captured two Lance physical layouts and {len(artifacts)} artifacts in {output}')