# /// script
# requires-python = ">=3.12,<3.13"
# dependencies = ["duckdb==1.4.4", "polars==1.44.1", "pyarrow==25.0.1"]
# ///
import argparse
import hashlib
import json
import platform
import tempfile
from datetime import datetime, timezone
from pathlib import Path

import duckdb
import polars as pl
import pyarrow as pa
import pyarrow.ipc as ipc


def digest_file(path):
    digest = hashlib.sha256()
    with path.open('rb') as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b''):
            digest.update(chunk)
    return digest.hexdigest()


def canonical(table):
    schema = pa.schema([(name, pa.int64()) for name in table.column_names])
    return table.cast(schema).combine_chunks()


def result_proof(table):
    table = canonical(table)
    sink = pa.BufferOutputStream()
    with ipc.new_stream(sink, table.schema) as writer:
        writer.write_table(table, max_chunksize=65536)
    return {'rows': table.num_rows, 'columns': table.column_names,
            'sha256': hashlib.sha256(sink.getvalue()).hexdigest(),
            'first': table.slice(0, 8).to_pylist(), 'last': table.slice(max(0, table.num_rows - 8)).to_pylist()}


def operators(profile, path='0'):
    output = []
    if 'operator_name' in profile:
        output.append({'path': path, 'name': profile['operator_name'], 'type': profile['operator_type'],
                       'outputRows': profile['operator_cardinality'], 'scannedRows': profile['operator_rows_scanned'],
                       'seconds': profile['operator_timing'], 'extra': profile['extra_info']})
    for index, child in enumerate(profile.get('children', [])):
        output.extend(operators(child, f'{path}.{index}'))
    return output


def metrics(profile):
    return {'peakBufferBytes': profile['system_peak_buffer_memory'],
            'peakTempBytes': profile['system_peak_temp_dir_size'],
            'rowsReturned': profile['rows_returned'],
            'cumulativeRowsScanned': profile['cumulative_rows_scanned'],
            'latencySeconds': profile['latency'], 'cpuSeconds': profile['cpu_time']}


parser = argparse.ArgumentParser()
parser.add_argument('--verify', action='store_true')
parser.add_argument('--output', type=Path, default=Path('explorer/static/engines/memory'))
args = parser.parse_args()
output = args.output.resolve()
source_sql = {
    'events': "SELECT i AS id, i % 500000 AS key, (i * 104729) % 1000000 AS score, repeat(md5(i::VARCHAR), 2) AS payload FROM range(1000000) t(i)",
    'lookup': "SELECT i AS key, i * 3 AS amount, repeat(md5(i::VARCHAR), 2) AS payload FROM range(500000) t(i)",
}
queries = {
    'aggregate-wide': 'SELECT key, sum(id)::BIGINT AS total, count(*)::BIGINT AS n FROM events GROUP BY key ORDER BY key',
    'aggregate-coarse': 'SELECT key % 4096 AS bucket, sum(id)::BIGINT AS total, count(*)::BIGINT AS n FROM events GROUP BY bucket ORDER BY bucket',
    'join': 'SELECT sum(length(e.payload || d.payload))::BIGINT AS checksum, count(*)::BIGINT AS n FROM events e JOIN lookup d USING (key)',
    'sort': 'SELECT id FROM events ORDER BY payload, id',
    'top-k': 'SELECT id FROM events ORDER BY payload, id LIMIT 10',
}
artifacts = {}
runs = []
with tempfile.TemporaryDirectory(prefix='columnar-memory-work-') as work:
    root = Path(work)
    db = root / 'input.duckdb'
    connection = duckdb.connect(str(db), config={'threads': 1})
    fixtures = {}
    for name, sql in source_sql.items():
        connection.execute(f'CREATE TABLE {name} AS {sql}')
        path = root / f'{name}.parquet'
        order = 'id' if name == 'events' else 'key'
        connection.execute(f"COPY (SELECT * FROM {name} ORDER BY {order}) TO ? (FORMAT PARQUET, COMPRESSION SNAPPY)", [str(path)])
        fixtures[name] = {'rows': 1000000 if name == 'events' else 500000, 'sql': sql, 'parquetBytes': path.stat().st_size, 'parquetSha256': digest_file(path)}
    events = pl.read_parquet(root / 'events.parquet')
    lookup = pl.read_parquet(root / 'lookup.parquet')
    assert events['id'].equals(pl.int_range(0, 1000000, eager=True).rename('id'))
    assert (events['key'] == events['id'] % 500000).all()
    assert (events['score'] == (events['id'] * 104729) % 1000000).all()
    for name, frame in [('events', events), ('lookup', lookup)]:
        arrow = connection.execute(f'SELECT * FROM {name}').fetch_arrow_table().combine_chunks()
        assert arrow.equals(frame.to_arrow().cast(arrow.schema).combine_chunks()), name
    connection.close()
    ordered = events.sort(['payload', 'id']).select('id')
    expected = {
        'aggregate-wide': events.group_by('key').agg(pl.col('id').sum().alias('total'), pl.len().alias('n')).sort('key'),
        'aggregate-coarse': events.with_columns((pl.col('key') % 4096).alias('bucket')).group_by('bucket').agg(pl.col('id').sum().alias('total'), pl.len().alias('n')).sort('bucket'),
        'join': events.join(lookup, on='key').select(pl.concat_str('payload', 'payload_right').str.len_bytes().cast(pl.Int64).sum().alias('checksum'), pl.len().alias('n')),
        'sort': ordered, 'top-k': ordered.head(10),
    }
    oracle = {name: result_proof(frame.to_arrow()) for name, frame in expected.items()}
    for name, sql in queries.items():
        for memory_limit in ['48MB', '256MB']:
            run_id = name + '-' + memory_limit
            profile_path = root / (run_id + '.json')
            connection = duckdb.connect(str(db), read_only=True, config={
                'threads': 1, 'memory_limit': memory_limit, 'preserve_insertion_order': False,
                'temp_directory': str(root / ('spill-' + run_id)), 'max_temp_directory_size': '1GB',
            })
            connection.execute("SET enable_profiling='json'")
            connection.execute('SET profiling_output = ?', [str(profile_path)])
            record = {'id': run_id, 'query': name, 'memoryLimit': memory_limit, 'sql': sql}
            try:
                table = connection.execute(sql).fetch_arrow_table()
                assert canonical(table).equals(canonical(expected[name].to_arrow())), run_id
                result = result_proof(table)
                assert result == oracle[name], run_id
                profile = json.loads(profile_path.read_text().replace(str(root), '<workspace>'))
                record.update(outcome='ok', result=result, metrics=metrics(profile), operators=operators(profile), profile=run_id + '.json')
                artifacts[run_id + '.json'] = (json.dumps(profile, indent=2) + '\n').encode()
            except duckdb.OutOfMemoryException as error:
                if name != 'aggregate-wide' or memory_limit != '48MB':
                    raise
                record.update(outcome='out-of-memory', error=str(error).replace(str(root), '<workspace>'))
            finally:
                connection.close()
            runs.append(record)
            print(run_id, record['outcome'], record.get('metrics', {}), flush=True)

proof = {'schemaVersion': 1, 'engine': 'DuckDB ' + duckdb.__version__, 'oracle': 'Polars ' + pl.__version__ + ' / PyArrow ' + pa.__version__,
         'recordedAt': datetime.now(timezone.utc).isoformat(), 'platform': {'system': platform.system(), 'architecture': platform.machine()},
         'configuration': {'threads': 1, 'preserveInsertionOrder': False, 'maxTempDirectorySize': '1GB', 'freshConnectionPerRun': True},
         'normalization': 'Only the temporary workspace path is replaced. Timings and native metrics are retained. Result digests are canonical uncompressed Arrow IPC streams with int64 columns and batches of 65536 rows.',
         'fixtures': fixtures, 'oracleResults': oracle, 'runs': runs,
         'artifacts': {name: {'bytes': len(data), 'sha256': hashlib.sha256(data).hexdigest()} for name, data in artifacts.items()}}
if args.verify:
    recorded = json.loads((output / 'proof.json').read_text())
    assert recorded['fixtures'] == fixtures
    assert recorded['oracleResults'] == oracle
    assert recorded['engine'] == proof['engine'] and recorded['oracle'] == proof['oracle']
    assert recorded['configuration'] == proof['configuration']
    assert [run['id'] for run in recorded['runs']] == [run['id'] for run in runs]
    assert set(recorded['artifacts']) == {run['profile'] for run in recorded['runs'] if run['outcome'] == 'ok'}
    for run in recorded['runs']:
        assert run['sql'] == queries[run['query']]
        assert run['memoryLimit'] in ['48MB', '256MB']
        if run['outcome'] == 'ok':
            assert run['result'] == oracle[run['query']]
            data = (output / run['profile']).read_bytes()
            assert recorded['artifacts'][run['profile']] == {'bytes': len(data), 'sha256': hashlib.sha256(data).hexdigest()}
            profile = json.loads(data)
            assert profile['query_name'] == run['sql']
            assert profile['rows_returned'] == run['result']['rows']
            assert run['metrics'] == metrics(profile) and run['operators'] == operators(profile)
        else:
            assert run['outcome'] == 'out-of-memory' and run['id'] == 'aggregate-wide-48MB'
            assert 'Out of Memory Error' in run['error']
    assert (output / 'reproduce.py').read_bytes() == Path(__file__).read_bytes()
else:
    artifacts['proof.json'] = (json.dumps(proof, indent=2) + '\n').encode()
    artifacts['reproduce.py'] = Path(__file__).read_bytes()
    output.mkdir(parents=True, exist_ok=True)
    for name, data in artifacts.items():
        with (output / name).open('xb') as stream:
            stream.write(data)
print('Verified every successful native result against complete independent Polars results. Runtime and spill quantities are observations, not portable performance thresholds.')