# /// script
# requires-python = ">=3.12,<3.13"
# dependencies = ["pyspark==3.5.6", "pyarrow==25.0.1", "duckdb==1.4.4", "fastavro==1.12.1"]
# ///
import argparse
import csv
import hashlib
import json
import os
import shutil
import subprocess
import sys
import tempfile
from datetime import datetime
from pathlib import Path
from urllib.parse import unquote, urlparse

import duckdb
import fastavro
import pyarrow as pa
import pyarrow.parquet as pq
from pyspark.sql import SparkSession

parser = argparse.ArgumentParser()
parser.add_argument('--verify', action='store_true')
parser.add_argument('--source', type=Path, default=Path('explorer/static/iceberg/weather/observations.csv'))
parser.add_argument('--output', type=Path, default=Path('explorer/static/iceberg/deletes'))
args = parser.parse_args()
output = args.output.resolve()
work = Path(tempfile.mkdtemp(prefix='iceberg-',dir='/private/tmp' if Path('/private/tmp').exists() else None))
with args.source.open() as stream:
    source = [{'observation_id':int(r['observation_id']),'station':r['station'],'observed_at':datetime.fromisoformat(r['observed_at']),'temperature_c':float(r['temperature_c'])} for r in csv.DictReader(stream) if int(r['observation_id'])<=18]
fields = [('observation_id',pa.int64(),'long'),('station',pa.string(),'string'),('observed_at',pa.timestamp('us'),'timestamp'),('temperature_c',pa.float64(),'double')]
arrow = pa.schema([pa.field(name,typ,metadata={'PARQUET:field_id':str(i+1)}) for i,(name,typ,_) in enumerate(fields)])
position_schema = pa.schema([pa.field('file_path',pa.string(),nullable=False,metadata={'PARQUET:field_id':'2147483546'}),pa.field('pos',pa.int64(),nullable=False,metadata={'PARQUET:field_id':'2147483545'})])


def norm(value):
    if isinstance(value,datetime): return value.isoformat(timespec='milliseconds')
    if isinstance(value,int) and abs(value)>9007199254740991: return str(value)
    if isinstance(value,bytes): return {'hex':value.hex()}
    if isinstance(value,dict): return {k:norm(v) for k,v in value.items()}
    if isinstance(value,(list,tuple)): return [norm(v) for v in value]
    return value


def path_of(uri):
    return Path(unquote(urlparse(uri).path)) if uri.startswith('file:') else Path(uri)


def replace(value,old,new):
    if isinstance(value,str): return value.replace(old,new)
    if isinstance(value,list): return [replace(v,old,new) for v in value]
    if isinstance(value,dict): return {k:replace(v,old,new) for k,v in value.items()}
    return value


def relocate():
    proof=json.loads((output/'proof.json').read_text())
    for artifact in proof['artifacts']:
        raw=(output/artifact['path']).read_bytes()
        assert len(raw)==artifact['size'] and hashlib.sha256(raw).hexdigest()==artifact['sha256']
    shutil.copytree(output/'tables',work/'tables')
    old,new=proof['captureRoot'],str(work)
    for file in (work/'tables').rglob('*.parquet'):
        data=pq.read_table(file)
        if data.schema.names==['file_path','pos']:
            data=data.set_column(0,data.schema.field(0),pa.array([v.replace(old,new) for v in data.column(0).to_pylist()]))
            pq.write_table(data,file,compression='NONE',use_dictionary=False,write_statistics=False)
    avros=[]
    for file in (work/'tables').rglob('*.avro'):
        with file.open('rb') as stream:
            reader=fastavro.reader(stream)
            avros.append((file,reader.writer_schema,reader.metadata,replace(list(reader),old,new)))
    for is_list in [False,True]:
        for file,schema,metadata,records in avros:
            if bool(records and 'manifest_path' in records[0]) != is_list: continue
            for record in records:
                if is_list: record['manifest_length']=path_of(record['manifest_path']).stat().st_size
                else: record['data_file']['file_size_in_bytes']=path_of(record['data_file']['file_path']).stat().st_size
            with file.open('wb') as stream:
                fastavro.writer(stream,schema,records,codec=metadata.get('avro.codec','null'),metadata={k:v for k,v in metadata.items() if k not in ['avro.schema','avro.codec']})
    for file in (work/'tables').rglob('*.json'):
        file.write_text(file.read_text().replace(old,new))
    return proof


os.environ['PYSPARK_PYTHON']=sys.executable
spark=(SparkSession.builder.master('local[2]').appName('Columnar Iceberg delete evidence')
    .config('spark.jars.packages','org.apache.iceberg:iceberg-spark-runtime-3.5_2.12:1.9.2')
    .config('spark.sql.extensions','org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions')
    .config('spark.sql.session.timeZone','UTC').config('spark.ui.enabled','false')
    .config('spark.sql.shuffle.partitions','2').getOrCreate())
spark.sparkContext.setLogLevel('ERROR')
j=spark._jvm
schema=j.org.apache.iceberg.SchemaParser.fromJson(json.dumps({'type':'struct','schema-id':0,'fields':[{'id':i+1,'name':name,'required':False,'type':typ} for i,(name,_,typ) in enumerate(fields)]}))
con=duckdb.connect();con.execute('INSTALL iceberg;LOAD iceberg;')


def duck_snapshot(metadata,sid,partitioned):
    script="""import duckdb,json,sys
con=duckdb.connect();con.execute('LOAD iceberg')
params=[sys.argv[1],int(sys.argv[2])]
rows=[];warm=[]
for station in (['North','South'] if sys.argv[3]=='True' else [None]):
 scope=' WHERE station=?' if station is not None else ''
 bound=params+[station] if station is not None else params
 rows.extend(con.execute('SELECT observation_id,station,observed_at,temperature_c FROM iceberg_scan(?,snapshot_from_id=?)'+scope,bound).fetchall())
 warm.extend(con.execute('SELECT observation_id,temperature_c,station FROM iceberg_scan(?,snapshot_from_id=?)'+scope+(' AND ' if scope else ' WHERE ')+'temperature_c>=20',bound).fetchall())
rows.sort(key=lambda r:(r[0],r[3]));warm.sort()
print(json.dumps({'rows':[[r[0],r[1],r[2].isoformat(timespec='milliseconds'),r[3]] for r in rows],'warm':[r[0] for r in warm]}))
"""
    return json.loads(subprocess.run([sys.executable,'-c',script,str(metadata),str(sid),str(partitioned)],check=True,stdout=subprocess.PIPE,text=True).stdout)


def inspect(table,path,label,expected):
    current=table.currentSnapshot()
    sid=current.snapshotId()
    actual=spark.read.format('iceberg').option('snapshot-id',str(sid)).load(path.as_uri()).orderBy('observation_id','temperature_c').collect()
    actual=[norm(row.asDict()) for row in actual]
    expected=sorted(norm(expected),key=lambda row:(row['observation_id'],row['temperature_c']))
    assert actual==expected,(label,actual,expected)
    metadata=table.operations().current().metadataFileLocation()
    oracle_result=duck_snapshot(metadata,sid,table.spec().isPartitioned())
    oracle=oracle_result['rows']
    oracle_rows=norm([dict(zip([f[0] for f in fields],row,strict=True)) for row in oracle])
    assert oracle_rows==expected,(label,oracle_rows,expected)
    warm=[row['observation_id'] for row in expected if row['temperature_c']>=20]
    assert oracle_result['warm']==warm
    assert [r.observation_id for r in spark.read.format('iceberg').option('snapshot-id',str(sid)).load(path.as_uri()).where('temperature_c>=20').select('observation_id').orderBy('observation_id').collect()]==warm
    projected=spark.read.format('iceberg').option('snapshot-id',str(sid)).load(path.as_uri()).select('temperature_c').collect()
    assert sorted(row.temperature_c for row in projected)==sorted(row['temperature_c'] for row in expected)
    entries=[]
    manifests=current.allManifests(table.io())
    for manifest in manifests:
        mp=path_of(str(manifest.path()))
        with mp.open('rb') as stream: raw=list(fastavro.reader(stream))
        for entry in raw:
            entry=norm(entry)
            entry['manifest']=str(mp.relative_to(work))
            entry['manifestSequence']=manifest.sequenceNumber()
            entry['effectiveSequence']=entry['sequence_number'] if entry['sequence_number'] is not None else manifest.sequenceNumber()
            entry['effectiveFileSequence']=entry['file_sequence_number'] if entry['file_sequence_number'] is not None else manifest.sequenceNumber()
            entries.append(entry)
    return {'label':label,'id':str(sid),'sequence':current.sequenceNumber(),'metadata':str(path_of(metadata).relative_to(work)),'manifestList':str(path_of(str(current.manifestListLocation())).relative_to(work)),'entries':entries,'rows':actual,'warmIds':warm}


if args.verify:
    proof=relocate()
    assert proof['sourceRows']==norm(source)
    for case in proof['cases']:
        path=work/'tables'/case['key']
        table=j.org.apache.iceberg.hadoop.HadoopTables(spark._jsc.hadoopConfiguration()).load(path.as_uri())
        for snap in case['snapshots']:
            sid=int(snap['id'])
            actual=spark.read.format('iceberg').option('snapshot-id',snap['id']).load(path.as_uri()).orderBy('observation_id','temperature_c').collect()
            assert norm([r.asDict() for r in actual])==snap['rows']
            metadata=str(work/snap['metadata'])
            oracle_result=duck_snapshot(metadata,sid,case['partitioned'])
            rows=oracle_result['rows']
            assert norm([dict(zip([f[0] for f in fields],r,strict=True)) for r in rows])==snap['rows']
            assert oracle_result['warm']==snap['warmIds']
            assert [r.observation_id for r in spark.read.format('iceberg').option('snapshot-id',snap['id']).load(path.as_uri()).where('temperature_c>=20').select('observation_id').orderBy('observation_id').collect()]==snap['warmIds']
            projected=spark.read.format('iceberg').option('snapshot-id',snap['id']).load(path.as_uri()).select('temperature_c').collect()
            assert sorted(r.temperature_c for r in projected)==sorted(r['temperature_c'] for r in snap['rows'])
    print('Verified shipped artifact hashes and every relocated snapshot/predicate/projection with Spark and DuckDB.')
else:
    assert not output.exists(),'Choose a fresh output directory; existing artifacts are never overwritten.'
    cases=[]
    for kind in ['position','equality','position-same','partition']:
        path=work/'tables'/kind;path.mkdir(parents=True)
        spec=j.org.apache.iceberg.PartitionSpec.builderFor(schema).identity('station').build() if kind=='partition' else j.org.apache.iceberg.PartitionSpec.unpartitioned()
        props=j.java.util.HashMap();props.put('format-version','2')
        table=j.org.apache.iceberg.hadoop.HadoopTables(spark._jsc.hadoopConfiguration()).create(schema,spec,props,path.as_uri())
        file_info=[]
        def write_data(name,rows,partition=None):
            file=path/(name+'.parquet');pq.write_table(pa.Table.from_pylist(rows,schema=arrow),file,compression='NONE',use_dictionary=False)
            builder=j.org.apache.iceberg.DataFiles.builder(spec).withPath(file.as_uri()).withFileSizeInBytes(file.stat().st_size).withRecordCount(len(rows))
            if partition is not None: builder=builder.withPartitionPath('station='+partition)
            file_info.append({'label':name,'path':str(file.relative_to(work)),'uri':file.as_uri(),'rows':norm(rows),'schema':[{'name':f.name,'id':int(f.metadata[b'PARQUET:field_id'])} for f in arrow]})
            return builder.build()
        def write_delete(name,positions=None,partition=None):
            file=path/(name+'.parquet')
            ds=position_schema if positions is not None else pa.schema([arrow.field(0)])
            values=positions if positions is not None else [{'observation_id':3}]
            pq.write_table(pa.Table.from_pylist(values,schema=ds),file,compression='NONE',use_dictionary=False,write_statistics=False)
            builder=j.org.apache.iceberg.FileMetadata.deleteFileBuilder(spec)
            if positions is not None: builder=builder.ofPositionDeletes()
            else:
                ids=spark.sparkContext._gateway.new_array(j.int,1);ids[0]=1;builder=builder.ofEqualityDeletes(ids)
            builder=builder.withPath(file.as_uri()).withFileSizeInBytes(file.stat().st_size).withRecordCount(len(values))
            if partition is not None: builder=builder.withPartitionPath('station='+partition)
            file_info.append({'label':name,'path':str(file.relative_to(work)),'uri':file.as_uri(),'rows':values,'schema':[{'name':f.name,'id':int(f.metadata[b'PARQUET:field_id'])} for f in ds]})
            return builder.build()
        first=[r for r in source if r['station']=='North'] if kind=='partition' else source[:6]
        rest=[r for r in source if r['station']=='South'] if kind=='partition' else source[6:]
        a=write_data('A',first,'North' if kind=='partition' else None)
        b=write_data('B',rest,'South' if kind=='partition' else None)
        table.newAppend().appendFile(a).appendFile(b).commit()
        snaps=[inspect(table,path,'Eighteen original observations',source)]
        replacement={**source[2],'temperature_c':30.0}
        if kind=='position':
            delete=write_delete('P',positions=[{'file_path':str(a.path()),'pos':2}])
            table.newRowDelta().addDeletes(delete).commit()
            survivors=[r for r in source if r['observation_id']!=3]
            snaps.append(inspect(table,path,'Delete A at position 2',survivors))
            table.newAppend().appendFile(write_data('D',[replacement])).commit()
            snaps.append(inspect(table,path,'Append a new observation 3',survivors+[replacement]))
        elif kind=='equality':
            table.newRowDelta().addDeletes(write_delete('E')).addRows(write_data('D',[replacement])).commit()
            snaps.append(inspect(table,path,'Delete key 3 and insert its replacement together',[r for r in source if r['observation_id']!=3]+[replacement]))
        elif kind=='position-same':
            data=write_data('D',[replacement])
            delete=write_delete('P',positions=[{'file_path':str(data.path()),'pos':0}])
            table.newRowDelta().addRows(data).addDeletes(delete).commit()
            snaps.append(inspect(table,path,'Insert D and delete its position 0 together',source))
        else:
            table.newRowDelta().addDeletes(write_delete('E-south',partition='South')).commit()
            snaps.append(inspect(table,path,'Delete key 3 scoped to South',source))
            table.newRowDelta().addDeletes(write_delete('E-north',partition='North')).commit()
            snaps.append(inspect(table,path,'Delete key 3 scoped to North',[r for r in source if r['observation_id']!=3]))
        cases.append({'key':kind,'partitioned':kind=='partition','spec':json.loads(j.org.apache.iceberg.PartitionSpecParser.toJson(spec)),'files':file_info,'snapshots':snaps})
    output.mkdir(parents=True)
    shutil.copytree(work/'tables',output/'tables',ignore=shutil.ignore_patterns('.*'))
    artifacts=[]
    for file in sorted(output.rglob('*')):
        if file.is_file():
            raw=file.read_bytes();artifacts.append({'path':str(file.relative_to(output)),'size':len(raw),'sha256':hashlib.sha256(raw).hexdigest()})
    proof={'writer':'PyArrow 25.0.1 files; Apache Iceberg Java 1.9.2 commits','readers':['Spark 3.5.6 / Iceberg 1.9.2','DuckDB 1.4.4 / iceberg extension'],'readerLimitations':['DuckDB 1.4.4 equality-delete narrow filter projection hits upstream issue 940. The oracle retains observation_id and temperature_c; Spark separately verifies narrow projections.', 'DuckDB unrestricted partitioned equality-delete scans removed a key in another partition in 1.4.4 and 1.5.2. The oracle combines explicit North and South scans; Spark verifies the unrestricted scan.'],'captureRoot':str(work),'sourceRows':norm(source),'cases':cases,'artifacts':artifacts}
    (output/'proof.json').write_text(json.dumps(proof,indent=2)+'\n')
    print('Generated and independently read',sum(len(c['snapshots']) for c in cases),'snapshots at',output)
spark.stop()