# /// script
# requires-python = ">=3.12,<3.13"
# dependencies = ["pyspark==3.5.6", "delta-spark==3.3.2", "pyarrow==25.0.1"]
# ///
import argparse
import hashlib
import json
import os
import shutil
import sys
import tempfile
import zipfile
from pathlib import Path

import pyarrow.parquet as pq
from delta import configure_spark_with_delta_pip
from py4j.protocol import Py4JJavaError
from pyspark.sql import SparkSession
from pyspark.sql import functions as sf

parser = argparse.ArgumentParser()
parser.add_argument('--source', type=Path, default=Path('explorer/static/delta/native/table'))
parser.add_argument('--output', type=Path, default=Path('explorer/static/delta/lifecycle'))
parser.add_argument('--verify', action='store_true')
args = parser.parse_args()
source = args.source.resolve()
output = args.output.resolve()
if not args.verify and output.exists():
    raise RuntimeError('Choose a fresh output directory; existing evidence is never replaced.')
work = Path(tempfile.mkdtemp(prefix='columnar-delta-lifecycle-'))
table = work / 'table'
capture = work / 'capture'
shutil.copytree(source, table)
capture.mkdir()
os.environ['PYSPARK_PYTHON'] = sys.executable
builder = (SparkSession.builder.master('local[2]').appName('Columnar isolated Delta lifecycle')
    .config('spark.sql.extensions', 'io.delta.sql.DeltaSparkSessionExtension')
    .config('spark.sql.catalog.spark_catalog', 'org.apache.spark.sql.delta.catalog.DeltaCatalog')
    .config('spark.sql.session.timeZone', 'UTC').config('spark.ui.enabled', 'false')
    .config('spark.sql.sources.parallelPartitionDiscovery.parallelism', '2')
    .config('spark.sql.shuffle.partitions', '2').config('spark.databricks.delta.snapshotPartitions', '2'))
spark = configure_spark_with_delta_pip(builder).getOrCreate()
spark.sparkContext.setLogLevel('ERROR')


def version(path):
    return max(int(file.stem) for file in (path / '_delta_log').glob('*.json'))


def active(path, selected):
    files = {}
    for index in range(selected + 1):
        actions = [json.loads(line) for line in (path / '_delta_log' / f'{index:020}.json').read_text().splitlines()]
        for action in actions:
            if 'remove' in action:
                files.pop(action['remove']['path'], None)
        for action in actions:
            if 'add' in action:
                files[action['add']['path']] = action['add']
    return files


def rows(path, selected):
    records = (spark.read.format('delta').option('versionAsOf', selected).load(str(path))
        .withColumn('observed_at', sf.date_format('observed_at', "yyyy-MM-dd'T'HH:mm:ss.SSSSSS"))
        .orderBy('observation_id').collect())
    return [row.asDict() for row in records]


def inventory(path):
    return [{'path': file.relative_to(path).as_posix(), 'size': file.stat().st_size,
             'sha256': hashlib.sha256(file.read_bytes()).hexdigest()}
            for file in sorted(path.rglob('*')) if file.is_file()]


def capture_step(label):
    current = version(table)
    selected = active(table, current)
    physical = {name: pq.read_table(table / name).column('observation_id').to_pylist() for name in selected}
    result = {'label': label, 'version': current, 'rows': rows(table, current), 'activeFiles': selected,
              'physicalIds': physical, 'inventory': inventory(table)}
    shutil.copytree(table, capture / label)
    return result


initial_version = version(table)
assert initial_version == 4
before = capture_step('before')
expected = [row for row in before['rows']]
expected_ids = [row['observation_id'] for row in expected]
assert expected_ids == [value for value in range(1, 19) if value not in [13, 16]]
assert sorted(value for ids in before['physicalIds'].values() for value in ids) == list(range(1, 19))
assert sum('deletionVector' in entry for entry in before['activeFiles'].values()) == 1

spark.sql(f'REORG TABLE delta.`{table}` APPLY (PURGE)').collect()
rewritten = capture_step('after-reorg')
assert rewritten['version'] == initial_version + 1
assert rewritten['rows'] == expected
assert all(not entry.get('deletionVector') for entry in rewritten['activeFiles'].values())
assert sorted(value for ids in rewritten['physicalIds'].values() for value in ids) == expected_ids
assert rows(table, initial_version) == expected
old_files = set(before['activeFiles']) - set(rewritten['activeFiles'])
assert old_files and all((table / name).exists() for name in old_files)
spark.sql(f'REORG TABLE delta.`{table}` APPLY (PURGE)').collect()
assert version(table) == rewritten['version']

spark.conf.set('spark.databricks.delta.retentionDurationCheck.enabled', 'false')
spark.sql(f'VACUUM delta.`{table}` RETAIN 0 HOURS').collect()
cleaned = capture_step('after-vacuum')
assert cleaned['rows'] == expected
cleanup_operations = []
for index in range(rewritten['version'] + 1, cleaned['version'] + 1):
    for line in (table / '_delta_log' / f'{index:020}.json').read_text().splitlines():
        action = json.loads(line)
        if 'commitInfo' in action:
            cleanup_operations.append(action['commitInfo']['operation'])
assert cleanup_operations == ['VACUUM START', 'VACUUM END']
assert all(not (table / name).exists() for name in old_files)
assert active(table, initial_version) == before['activeFiles']
historical_error = None
try:
    rows(table, initial_version)
except Py4JJavaError as error:
    message = str(error.java_exception)
    assert 'FileNotFoundException' in message or 'FILE_NOT_EXIST' in message, message
    historical_error = message.replace(str(work), '<temporary>')[:4000]
assert historical_error is not None, 'The historical read must fail after its required data file is vacuumed.'

steps = [before, rewritten, cleaned]
semantic = {'logicalIds': expected_ids, 'versions': [step['version'] for step in steps],
            'activePhysicalCounts': [sum(len(ids) for ids in step['physicalIds'].values()) for step in steps],
            'activeVectorCounts': [sum(bool(entry.get('deletionVector')) for entry in step['activeFiles'].values()) for step in steps],
            'currentRowsUnchanged': all(step['rows'] == expected for step in steps),
            'cleanupOperations': cleanup_operations,
            'reorgIdempotent': True, 'oldMembershipReconstructible': True, 'historicalReadFails': True}
proof = {'engine': 'Spark 3.5.6 / Delta Spark 3.3.2', 'scope': 'Isolated temporary copy of the native deletion-vector table; no source or production data is cleaned.',
         'sourceVersion': initial_version, 'operations': ['REORG TABLE APPLY (PURGE)', 'REORG TABLE APPLY (PURGE) again', 'VACUUM RETAIN 0 HOURS'],
         'retention': 'Zero retention with the duration check disabled only for this disposable experiment. It is not a production recommendation.',
         'steps': steps, 'removedHistoricalFiles': sorted(old_files), 'historicalError': historical_error,
         'semantic': semantic, 'artifacts': inventory(capture)}

if args.verify:
    recorded = json.loads((output / 'proof.json').read_text())
    assert recorded['semantic'] == semantic
    assert recorded['sourceVersion'] == initial_version
    assert recorded['removedHistoricalFiles'] == sorted(old_files)
    assert [step['label'] for step in recorded['steps']] == ['before', 'after-reorg', 'after-vacuum']
    for artifact in recorded['artifacts']:
        payload = (output / artifact['path']).read_bytes()
        assert len(payload) == artifact['size']
        assert hashlib.sha256(payload).hexdigest() == artifact['sha256']
    for step in recorded['steps']:
        path = output / step['label']
        assert inventory(path) == step['inventory']
        assert active(path, step['version']) == step['activeFiles']
        assert rows(path, step['version']) == step['rows']
        assert {name: pq.read_table(path / name).column('observation_id').to_pylist()
                for name in step['activeFiles']} == step['physicalIds']
    assert active(output / 'after-vacuum', initial_version) == recorded['steps'][0]['activeFiles']
    assert all((output / 'after-reorg' / name).exists()
               and not (output / 'after-vacuum' / name).exists() for name in old_files)
    assert (output / 'reproduce.py').read_bytes() == Path(__file__).read_bytes()
else:
    (capture / 'proof.json').write_text(json.dumps(proof, indent=2) + '\n')
    shutil.copy2(Path(__file__), capture / 'reproduce.py')
    shutil.copytree(capture, output)
published = recorded if args.verify else proof
for artifact in published['artifacts']:
    if any(part.startswith('.') for part in Path(artifact['path']).parts):
        download = output / 'downloads' / (artifact['sha256'] + '.bin')
        if args.verify:
            assert download.read_bytes() == (output / artifact['path']).read_bytes()
        else:
            download.parent.mkdir(exist_ok=True)
            download.write_bytes((output / artifact['path']).read_bytes())
archive = output.with_name(output.name + '.zip')
archive_names = sorted([entry['path'] for entry in proof['artifacts']] + ['proof.json', 'reproduce.py'])
if args.verify:
    archive_names = sorted([entry['path'] for entry in recorded['artifacts']] + ['proof.json', 'reproduce.py'])
    with zipfile.ZipFile(archive) as bundle:
        assert sorted(bundle.namelist()) == archive_names
        for name in archive_names:
            assert bundle.read(name) == (output / name).read_bytes()
else:
    with zipfile.ZipFile(archive, 'x', compression=zipfile.ZIP_DEFLATED, compresslevel=9) as bundle:
        for name in archive_names:
            info = zipfile.ZipInfo(name, date_time=(1980, 1, 1, 0, 0, 0))
            info.compress_type = zipfile.ZIP_DEFLATED
            bundle.writestr(info, (output / name).read_bytes())
spark.stop()
print('Verified native DV materialization, idempotent REORG, retained historical rows, physical VACUUM and failed historical read.')