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

import lance
import numpy as np
import pyarrow as pa
from lance.file import LanceFileReader


POINTS = [[5.125, 0.5], [510.125, 0.5], [1024.125, 0.5], [1030.125, 0.5]]
STAGES = ['indexed', 'append', 'delete', 'compact', 'reindex']
COMMANDS = {
    'indexed': "write_dataset(first_1024, max_rows_per_file=512, data_storage_version='2.1', enable_stable_row_ids=False); create_index('vector', 'IVF_FLAT', num_partitions=2, ivf_centroids=[[256, 0], [768, 0]])",
    'append': "write_dataset(last_8, dataset, mode='append', data_storage_version='2.1')",
    'delete': "dataset.delete('id IN (5, 1024)')",
    'compact': 'dataset.optimize.compact_files(target_rows_per_fragment=2048, num_threads=1)',
    'reindex': 'dataset.optimize.optimize_indices()',
}


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


def canonical(value):
    return json.dumps(value, sort_keys=True, separators=(',', ':')).encode()


def expected_ids(stage):
    ids = np.arange(1024 if stage == 'indexed' else 1032, dtype=np.int64)
    return ids if stage in ['indexed', 'append'] else ids[~np.isin(ids, [5, 1024])]


def addresses(rows):
    return [{**row, '_rowid': str(row['_rowid']), '_rowaddr': str(row['_rowaddr']),
             'fragment': str(row['_rowaddr'] >> 32), 'offset': row['_rowaddr'] & 0xffffffff}
            for row in rows]


def capture(dataset, stage, root):
    ids = expected_ids(stage)
    rows = sorted(dataset.to_table(with_row_id=True, with_row_address=True).to_pylist(), key=lambda row: row['id'])
    np.testing.assert_array_equal([row['id'] for row in rows], ids)
    np.testing.assert_array_equal([row['vector'] for row in rows], np.column_stack([ids, np.zeros(len(ids))]))
    assert all(row['_rowid'] == row['_rowaddr'] for row in rows)
    fragments = [fragment.metadata.to_json() for fragment in dataset.get_fragments()]
    for fragment in fragments:
        assert fragment['row_id_meta'] is None and not fragment['overlays']
        assert len(fragment['files']) == 1
        reference = fragment['files'][0]
        assert reference['fields'] == [0, 1] and reference['column_indices'] == [0, 1]
        assert (reference['file_major_version'], reference['file_minor_version']) == (2, 1)
        original = LanceFileReader(str(root / 'data' / reference['path'])).read_all().to_table()
        assert original.num_rows == fragment['physical_rows']
        selected = [row for row in rows if row['_rowaddr'] >> 32 == fragment['id']]
        offsets = [row['_rowaddr'] & 0xffffffff for row in selected]
        assert all(offset < original.num_rows for offset in offsets)
        assert original.take(pa.array(offsets, type=pa.int64())).to_pylist() == [{'id': row['id'], 'vector': row['vector']} for row in selected]
    assert {row['_rowaddr'] >> 32 for row in rows} <= {fragment['id'] for fragment in fragments}
    indices = [{**index, 'fragment_ids': sorted(index['fragment_ids'])} for index in dataset.list_indices()]
    assert len(indices) == 1 and indices[0]['type'] == 'IVF_FLAT'
    stats = dataset.index_statistics(indices[0]['name'])
    stats = json.loads(json.dumps(stats).replace(str(root), '<dataset>'))
    assert all(0 < index['num_partitions'] <= 2 for index in stats['indices'])
    assert stats['num_indexed_rows'] + stats['num_unindexed_rows'] == len(rows)
    covered = set(indices[0]['fragment_ids'])
    assert sum(row['_rowaddr'] >> 32 in covered for row in rows) == stats['num_indexed_rows']
    queries = []
    for point in POINTS:
        distances = (ids.astype(np.float64) - point[0]) ** 2 + point[1] ** 2
        order = np.argsort(distances)[:3]
        assert len(set(distances[order])) == 3
        scanner = dataset.scanner(nearest={'column': 'vector', 'q': point, 'k': 3, 'nprobes': 2}, with_row_id=True, with_row_address=True)
        actual = scanner.to_table().to_pylist()
        np.testing.assert_array_equal([row['id'] for row in actual], ids[order])
        np.testing.assert_array_equal([row['_distance'] for row in actual], distances[order])
        lookup = {row['id']: row for row in rows}
        assert all({key: value for key, value in row.items() if key != '_distance'} == lookup[row['id']] for row in actual)
        plan = scanner.explain_plan().replace(str(root).lstrip('/'), '<dataset>').replace(str(root), '<dataset>')
        queries.append({'point': point, 'rows': addresses(actual), 'exactIds': ids[order].tolist(), 'exactDistances': distances[order].tolist(), 'plan': plan})
    manifest = f'_versions/{2 ** 64 - 1 - dataset.version}.manifest'
    assert (root / manifest).is_file()
    return {'operation': stage, 'command': COMMANDS[stage], 'version': str(dataset.version),
            'manifest': manifest, 'rowCount': len(rows), 'fullRowsSha256': digest(canonical(addresses(rows))),
            'logicalRowsSha256': digest(canonical([{'id': row['id'], 'vector': row['vector']} for row in rows])),
            'traced': addresses([row for row in rows if row['id'] in [5, 6, 1023, 1024, 1025]]),
            'fragments': fragments, 'indices': indices, 'statistics': stats, 'queries': queries}


def verify_transitions(steps):
    indexed, appended, deleted, compacted, reindexed = steps
    assert [step['rowCount'] for step in steps] == [1024, 1032, 1030, 1030, 1030]
    assert [step['statistics']['num_unindexed_rows'] for step in steps] == [0, 8, 7, 7, 0]
    assert [len(step['fragments']) for step in steps] == [2, 3, 3, 2, 2]
    assert deleted['logicalRowsSha256'] == compacted['logicalRowsSha256'] == reindexed['logicalRowsSha256']
    for identity in [6, 1023, 1025]:
        before = next(row for row in deleted['traced'] if row['id'] == identity)
        after = next(row for row in compacted['traced'] if row['id'] == identity)
        assert before['_rowaddr'] != after['_rowaddr'] and before['vector'] == after['vector']
    assert compacted['fullRowsSha256'] == reindexed['fullRowsSha256']
    assert any(fragment['deletion_file'] for fragment in deleted['fragments'])
    assert all(fragment['deletion_file'] is None for fragment in compacted['fragments'])
    assert 'UnionExec' in appended['queries'][2]['plan'] and 'LanceScan' in appended['queries'][2]['plan']
    assert 'ANN' in reindexed['queries'][2]['plan'] and 'UnionExec' not in reindexed['queries'][2]['plan']


def verify_originals(output, proof):
    assert proof['schemaVersion'] == 1
    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('dataset/')}
        for path in archive.namelist():
            assert archive.read(path) == (output / path).read_bytes(), path
    for step in proof['steps']:
        retained = lance.dataset(output / 'dataset', version=int(step['version']))
        actual = capture(retained, step['operation'], output / 'dataset')
        assert actual == step, step['operation']
    verify_transitions(proof['steps'])


parser = argparse.ArgumentParser()
parser.add_argument('--verify', action='store_true')
parser.add_argument('--output', type=Path, default=Path('explorer/static/lance/lifecycle'))
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 lifecycle evidence before replacing it.')
with tempfile.TemporaryDirectory(prefix='columnar-lance-history-') as work:
    root = Path(work) / 'dataset'
    values = np.arange(1032, dtype=np.float32)
    vectors = np.column_stack([values, np.zeros(1032, dtype=np.float32)])
    original = pa.table({'id': pa.array(np.arange(1032), type=pa.int64()), 'vector': pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel()), 2)})
    dataset = lance.write_dataset(original.slice(0, 1024), root, max_rows_per_file=512, data_storage_version='2.1', enable_stable_row_ids=False)
    dataset.create_index('vector', 'IVF_FLAT', num_partitions=2, ivf_centroids=np.array([[256, 0], [768, 0]], dtype=np.float32))
    steps = []
    for stage in STAGES:
        if stage == 'append':
            dataset = lance.write_dataset(original.slice(1024), dataset, mode='append', data_storage_version='2.1')
        elif stage == 'delete':
            dataset.delete('id IN (5, 1024)')
        elif stage == 'compact':
            dataset.optimize.compact_files(target_rows_per_fragment=2048, num_threads=1)
        elif stage == 'reindex':
            dataset.optimize.optimize_indices()
        steps.append(capture(dataset, stage, root))
    verify_transitions(steps)
    if args.verify:
        proof = json.loads((output / 'proof.json').read_text())
        verify_originals(output, proof)
        for fresh, recorded in zip(steps, proof['steps'], strict=True):
            for field in ['operation', 'command', 'version', 'rowCount', 'fullRowsSha256', 'logicalRowsSha256', 'traced']:
                assert fresh[field] == recorded[field], (fresh['operation'], field)
            for left, right in zip(fresh['queries'], recorded['queries'], strict=True):
                for field in ['point', 'rows', 'exactIds', 'exactDistances']:
                    assert left[field] == right[field]
        assert (output / 'reproduce.py').read_bytes() == Path(__file__).read_bytes()
        print('Verified five retained Lance versions, every logical row and physical address, twenty exact NumPy queries, index coverage, original hashes and fresh lifecycle reproduction.')
    else:
        shutil.copytree(root, output / 'dataset')
        with zipfile.ZipFile(output / 'originals.zip', 'w', compression=zipfile.ZIP_DEFLATED) as archive:
            for path in sorted((output / 'dataset').rglob('*')):
                if path.is_file():
                    archive.write(path, path.relative_to(output))
        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', 'oracle': 'NumPy 2.4.3; source IDs and vectors constructed independently',
                 'format': '2.1', 'stableRowIds': False, 'points': POINTS, 'k': 3, 'nprobes': 2, 'steps': steps, 'artifacts': artifacts}
        (output / 'proof.json').write_text(json.dumps(proof, indent=2) + '\n')
        shutil.copyfile(__file__, output / 'reproduce.py')
        (output / 'README.md').write_text('''# Lance index lifecycle\n\nFive native versions retain original manifests, index files, data files and deletion files.\nThe application IDs and two-dimensional vectors are synthetic. NumPy checks twenty full\ntop-three answers; all logical rows and native physical addresses are checked against\noriginal fragment data. Stable row IDs are disabled. SDK manifests and index statistics\nare native captures, not an independent generic table decoder.\n\nSave reproduce.py and proof.json beside the extracted originals.zip, then run:\n\n```sh\nuv run reproduce.py --verify --output .\n```\n\nThe verifier reopens every preserved version and also builds a fresh native lifecycle.\nOriginal artifacts remain hash-checked; fresh UUIDs, timestamps and encoded index bytes\nmay differ. Compaction preserves separate indexed and unindexed replacement fragments.\nThe query probes all partitions in this small IVF_FLAT fixture; it is not an ANN speed\nbenchmark or a general exactness claim. No files have been vacuumed.\n''')
        verify_originals(output, proof)
        print(f'Captured five Lance versions and {len(artifacts)} original artifacts in {output}')