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

import fastavro
import pyarrow as pa
import pyarrow.parquet as pq
from py4j.protocol import Py4JJavaError
from pyspark.sql import SparkSession
from pyspark.errors import IllegalArgumentException

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/lifecycle'))
args=parser.parse_args()
work=Path(tempfile.mkdtemp(prefix='iceberg-lifecycle-'))
evidence=work/'evidence';evidence.mkdir()
output=args.output.resolve()
with args.source.open() as stream:
    source=[{'observation_id':int(row['observation_id']),'station':row['station'],'observed_at':datetime.fromisoformat(row['observed_at']),'temperature_c':float(row['temperature_c'])} for row in csv.DictReader(stream) if int(row['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)])
os.environ['PYSPARK_PYTHON']=sys.executable
spark=(SparkSession.builder.master('local[2]').appName('Columnar Iceberg lifecycle')
    .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
runtime=j.java.lang.Thread.currentThread().getContextClassLoader().loadClass('org.apache.iceberg.TableMetadata')
jar=Path(unquote(urlparse(str(runtime.getProtectionDomain().getCodeSource().getLocation())).path))
java=Path(os.environ['JAVA_HOME'])/'bin'
subprocess.run([str(java/'javac'),'-Xlint:deprecation','-cp',str(jar),'-d',str(work),str(Path(__file__).with_name('NativeVectors.java'))],check=True)
urls=spark.sparkContext._gateway.new_array(j.java.net.URL,1);urls[0]=j.java.net.URI(work.as_uri()+'/').toURL()
loader=j.java.net.URLClassLoader(urls,j.java.lang.Thread.currentThread().getContextClassLoader())
j.java.lang.Thread.currentThread().setContextClassLoader(loader)
helper=j.java.lang.Class.forName('NativeVectors',True,loader)
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)]}))
tables=j.org.apache.iceberg.hadoop.HadoopTables(spark._jsc.hadoopConfiguration())


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 {key:norm(item) for key,item in value.items()}
    if isinstance(value,(list,tuple)): return [norm(item) for item in value]
    return value


def path_of(uri):
    return Path(unquote(urlparse(str(uri)).path))


def save(table):
    directory=path_of(table.location())
    for file in directory.rglob('*'):
        if file.is_file() and not file.name.startswith('.'):
            dest=evidence/file.relative_to(work);dest.parent.mkdir(parents=True,exist_ok=True);shutil.copy2(file,dest)


def create(key,version=2,partitioned=False):
    spec=j.org.apache.iceberg.PartitionSpec.builderFor(schema).identity('station').build() if partitioned else j.org.apache.iceberg.PartitionSpec.unpartitioned()
    props=j.java.util.HashMap();props.put('format-version',str(version))
    return tables.create(schema,spec,props,(work/key).as_uri())


def data(table,name,rows,spec=None,partition=None):
    spec=spec if spec is not None else table.spec()
    file=path_of(table.location())/(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)
    metrics=j.org.apache.iceberg.parquet.ParquetUtil.fileMetrics(table.io().newInputFile(file.as_uri()),j.org.apache.iceberg.MetricsConfig.forTable(table))
    return builder.withMetrics(metrics).build()


def equality(table,name,rows,ids,spec=None,partition=None):
    spec=spec if spec is not None else table.spec()
    file=path_of(table.location())/(name+'.parquet')
    pq.write_table(pa.Table.from_pylist(rows,schema=pa.schema([arrow.field(i-1) for i in ids])),file,compression='NONE',use_dictionary=False)
    keys=spark.sparkContext._gateway.new_array(j.int,len(ids))
    for i,key in enumerate(ids): keys[i]=key
    builder=j.org.apache.iceberg.FileMetadata.deleteFileBuilder(spec).ofEqualityDeletes(keys).withPath(file.as_uri()).withFileSizeInBytes(file.stat().st_size).withRecordCount(len(rows))
    if partition is not None: builder=builder.withPartitionPath('station='+partition)
    metrics=j.org.apache.iceberg.parquet.ParquetUtil.fileMetrics(table.io().newInputFile(file.as_uri()),j.org.apache.iceberg.MetricsConfig.forTable(table))
    return builder.withMetrics(metrics).build()


def capture(table,label,expected):
    table.refresh();snap=table.currentSnapshot();sid=snap.snapshotId()
    frame=spark.read.format('iceberg').option('snapshot-id',str(sid)).load(table.location())
    rows=norm([row.asDict() for row in frame.orderBy('observation_id','station','temperature_c').collect()])
    expected=sorted(norm(expected),key=lambda r:(r['observation_id'],r['station'] or '',r['temperature_c']))
    assert rows==expected,(label,rows,expected)
    projected=sorted(row.temperature_c for row in frame.select('temperature_c').collect())
    assert projected==sorted(row['temperature_c'] for row in expected)
    files=[]
    tasks=table.newScan().useSnapshot(sid).planFiles()
    iterator=tasks.iterator()
    while iterator.hasNext():
        task=iterator.next()
        files.append({'path':str(path_of(task.file().location()).relative_to(work)),'records':task.file().recordCount(),'spec':task.spec().specId(),'deletes':[{'path':str(path_of(d.location()).relative_to(work)),'records':d.recordCount(),'content':str(d.content()),'spec':d.specId(),'offset':d.contentOffset(),'length':d.contentSizeInBytes(),'referencedDataFile':d.referencedDataFile(),'equalityIds':list(d.equalityFieldIds()) if d.equalityFieldIds() is not None else []} for d in task.deletes()]})
    tasks.close()
    entries=[]
    for manifest in snap.allManifests(table.io()):
        with path_of(manifest.path()).open('rb') as stream:
            entries.extend([{'manifest':str(path_of(manifest.path()).relative_to(work)),'manifestSequence':manifest.sequenceNumber(),'spec':manifest.partitionSpecId(),**norm(entry)} for entry in fastavro.reader(stream)])
    save(table)
    metadata_path=path_of(table.operations().current().metadataFileLocation())
    metadata=json.loads(metadata_path.read_text())
    return {'label':label,'id':str(sid),'sequence':snap.sequenceNumber(),'metadata':str(metadata_path.relative_to(work)),'rows':rows,'files':files,'entries':entries,'refs':{name:str(ref['snapshot-id']) for name,ref in metadata.get('refs',{}).items()},'retainedSnapshots':[str(item['snapshot-id']) for item in metadata['snapshots']],'physicalFiles':sorted(str(file.relative_to(work)) for file in path_of(table.location()).rglob('*.parquet'))}


def rejection(action,expected_class):
    try: action()
    except Py4JJavaError as error:
        name=error.java_exception.getClass().getName()
        assert name==expected_class,(name,str(error.java_exception))
        return {'class':name,'message':str(error.java_exception)}
    except IllegalArgumentException as error:
        assert expected_class=='java.lang.IllegalArgumentException'
        return {'class':expected_class,'message':str(error)}
    raise AssertionError('Expected the native operation to reject')


cases=[]
for kind in ['conflict','rewrite-append','append-append']:
    table=create(kind);a=data(table,'A',source);table.newAppend().appendFile(a).commit()
    base=capture(table,'Both writers read the same snapshot',source);sid=table.currentSnapshot().snapshotId()
    other=tables.load(table.location());stale=other.currentSnapshot().snapshotId();assert stale==sid
    first=data(table,'writer-A',[row for row in source if row['observation_id']!=3]) if kind!='append-append' else data(table,'writer-A',[source[2]])
    second=data(table,'writer-B',[row for row in source if row['observation_id']!=4]) if kind=='conflict' else data(table,'writer-B',[source[3]])
    pending=(other.newOverwrite().deleteFile(a).addFile(second).validateFromSnapshot(sid).validateNoConflictingData().validateNoConflictingDeletes()) if kind=='conflict' else other.newAppend().appendFile(second)
    if kind=='append-append': table.newAppend().appendFile(first).commit();after_rows=source+[source[2]]
    else: table.newOverwrite().deleteFile(a).addFile(first).commit();after_rows=[row for row in source if row['observation_id']!=3]
    after=capture(table,'Writer A commits',after_rows)
    error=rejection(lambda:pending.commit(),'org.apache.iceberg.exceptions.ValidationException') if kind=='conflict' else None
    if error:
        rejected=capture(table,'B rejects; A remains committed',after_rows)
        recomputed=[row for row in after_rows if row['observation_id']!=4]
        table.newOverwrite().deleteFile(first).addFile(data(table,'fresh-B',recomputed)).commit()
        stages=[base,after,rejected,capture(table,'B recomputes from current rows',recomputed)]
    else:
        pending.commit();stages=[base,after,capture(table,'Stale B append rebases and commits',after_rows+[source[3]])]
    cases.append({'key':kind,'stages':stages,'error':error,'staleSnapshot':str(stale),'stagedRows':norm(pq.read_table(path_of(second.location())).to_pylist())})
    print('Verified',kind,flush=True)

table=create('keys');rows=[{**row,'station':None} if row['observation_id'] in [3,4] else row for row in source]
rows.append({**source[2],'station':'South','temperature_c':30.0})
a=data(table,'A',rows);table.newAppend().appendFile(a).commit();stages=[capture(table,'A duplicate ID and two null stations',rows)]
delete=equality(table,'key-3-null',[{'observation_id':3,'station':None}],[1,2]);table.newRowDelta().addDeletes(delete).commit()
survivors=[row for row in rows if not(row['observation_id']==3 and row['station'] is None)]
stages.append(capture(table,'Composite key uses AND and null-safe matching',survivors))
cases.append({'key':'keys','stages':stages})
print('Verified composite/null keys',flush=True)

table=create('cross-spec');old=table.spec();a=data(table,'old-A',source[:6]);table.newAppend().appendFile(a).commit()
table.updateSpec().addField('station').commit()
b=data(table,'North',source[6::2],partition='North');c=data(table,'South',source[7::2],partition='South')
table.newAppend().appendFile(b).appendFile(c).commit();stages=[capture(table,'Old unpartitioned and new station partitions',source)]
delete=equality(table,'global',[{'observation_id':3},{'observation_id':9},{'observation_id':10}],[1],old)
table.newRowDelta().addDeletes(delete).commit();stages.append(capture(table,'Global equality delete reaches both specs',[r for r in source if r['observation_id'] not in [3,9,10]]))
cases.append({'key':'cross-spec','stages':stages});print('Verified cross-spec delete',flush=True)

table=create('compaction');a=data(table,'A',source[:6]);b=data(table,'B',source[6:]);table.newAppend().appendFile(a).appendFile(b).commit()
delete=equality(table,'E',[{'observation_id':3}],[1]);table.newRowDelta().addDeletes(delete).commit()
survivors=[r for r in source if r['observation_id']!=3];stages=[capture(table,'Read A plus its equality delete',survivors)]
table.newRewrite().validateFromSnapshot(table.currentSnapshot().snapshotId()).deleteFile(a).deleteFile(delete).addFile(data(table,'compacted',[r for r in survivors if r['observation_id']<=6])).commit()
stages.append(capture(table,'Rewrite survivors and retire the delete',survivors));assert all(not file['deletes'] for file in stages[-1]['files'])
assert [len(file['deletes']) for file in sorted(stages[0]['files'],key=lambda file:file['path'])]==[1,0]
cases.append({'key':'compaction','stages':stages});print('Verified delete compaction',flush=True)

table=create('cleanup');a=data(table,'A',source);table.newAppend().appendFile(a).commit();base=capture(table,'Original snapshot',source);sid=table.currentSnapshot().snapshotId()
table.manageSnapshots().createBranch('audit',sid).commit()
survivors=[r for r in source if r['observation_id']!=3];replacement=data(table,'replacement',survivors)
orphan=data(table,'never-committed',source);orphan_path=path_of(orphan.location());old_time=time.time()-10*86400;os.utime(orphan_path,(old_time,old_time))
table.newOverwrite().deleteFile(a).addFile(replacement).commit();latest=capture(table,'Main replaced A; audit still pins A',survivors)
error=rejection(lambda:table.expireSnapshots().expireSnapshotId(sid).commit(),'java.lang.IllegalArgumentException')
assert path_of(a.location()).exists()
table.manageSnapshots().removeBranch('audit').commit();table.expireSnapshots().expireSnapshotId(sid).commit()
assert not path_of(a.location()).exists();assert orphan_path.exists();assert path_of(replacement.location()).exists()
expired=capture(table,'Expiration removes unreferenced historical data',survivors)
unavailable=rejection(lambda:table.newScan().useSnapshot(sid),'java.lang.IllegalArgumentException')
result=j.org.apache.iceberg.spark.actions.SparkActions.get(spark._jsparkSession).deleteOrphanFiles(table).olderThan(int((time.time()-7*86400)*1000)).execute()
removed=[str(path_of(uri).relative_to(work)) for uri in result.orphanFileLocations()]
assert removed==['cleanup/never-committed.parquet'],removed
assert not orphan_path.exists();assert path_of(replacement.location()).exists()
final=capture(table,'Orphan cleanup removes only old unpublished output',survivors)
cases.append({'key':'cleanup','stages':[base,latest,expired,final],'error':error,'expiredRead':unavailable,'removedOrphans':removed,'filesAfterExpiration':['replacement.parquet','never-committed.parquet'],'filesAfterCleanup':['replacement.parquet']})
print('Verified expiration and physical orphan cleanup',flush=True)

table=create('vectors',3);a=data(table,'A',source);table.newAppend().appendFile(a).commit();stages=[capture(table,'v3 data before deletion',source)];previous=None;vectors=[]
for position in [2,15]:
    vector=j.NativeVectors.write(table,a,previous,position)
    update=table.newRowDelta().validateFromSnapshot(table.currentSnapshot().snapshotId()).addDeletes(vector)
    if previous is not None: update=update.removeDeletes(previous)
    update.commit()
    positions=list(j.NativeVectors.positions(table,vector));expected_positions=[2] if previous is None else [2,15];assert positions==expected_positions
    stages.append(capture(table,'DV masks positions '+str(positions),[row for i,row in enumerate(source) if i not in positions]))
    vectors.append({'path':str(path_of(vector.location()).relative_to(work)),'offset':vector.contentOffset(),'length':vector.contentSizeInBytes(),'positions':positions,'referencedDataFile':str(path_of(a.location()).relative_to(work)),'referencedDataUri':a.location()})
    previous=vector
cases.append({'key':'vectors','stages':stages,'vectors':vectors});print('Verified cumulative v3 deletion vectors',flush=True)

artifacts=[]
for file in sorted(evidence.rglob('*')):
    if file.is_file():
        raw=file.read_bytes();artifacts.append({'path':str(file.relative_to(evidence)),'size':len(raw),'sha256':hashlib.sha256(raw).hexdigest()})
proof={'writer':'Apache Iceberg Java 1.9.2 / PyArrow 25.0.1','reader':'Spark 3.5.6 / Iceberg 1.9.2','captureRoot':str(work),'sourceRows':norm(source),'cases':cases,'artifacts':artifacts}
if args.verify:
    original=json.loads((output/'proof.json').read_text());assert original['sourceRows']==proof['sourceRows']
    for artifact in original['artifacts']:
        raw=(output/artifact['path']).read_bytes();assert len(raw)==artifact['size'] and hashlib.sha256(raw).hexdigest()==artifact['sha256']
    for old,new in zip(original['cases'],cases,strict=True):
        assert old['key']==new['key']
        assert [s['rows'] for s in old['stages']]==[s['rows'] for s in new['stages']]
        assert (old.get('error') or {}).get('class')==(new.get('error') or {}).get('class')
        assert (old.get('expiredRead') or {}).get('class')==(new.get('expiredRead') or {}).get('class')
        assert old.get('stagedRows')==new.get('stagedRows')
        assert old.get('removedOrphans')==new.get('removedOrphans')
        assert [v['positions'] for v in old.get('vectors',[])]==[v['positions'] for v in new.get('vectors',[])]
        def signature(case):
            ids=list(dict.fromkeys(sid for stage in case['stages'] for sid in stage['retainedSnapshots']))
            result=[]
            for stage in case['stages']:
                tasks=[]
                for file in stage['files']:
                    deletes=sorted((d['content'],d['records'],d['spec'],d['offset'],d['length'],tuple(d['equalityIds'])) for d in file['deletes'])
                    tasks.append((file['path'],file['records'],file['spec'],deletes))
                result.append((stage['sequence'],sorted(tasks),stage['physicalFiles'],{name:ids.index(sid) for name,sid in stage['refs'].items()},[ids.index(sid) for sid in stage['retainedSnapshots']]))
            return result
        assert signature(old)==signature(new)
    print('Verified original hashes and independently regenerated every native outcome.')
else:
    assert not output.exists(),'Choose a fresh output directory.'
    shutil.copytree(evidence,output);(output/'proof.json').write_text(json.dumps(proof,indent=2)+'\n')
    print('Generated native lifecycle evidence:',output)
spark.stop()