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

import duckdb
import pyarrow.parquet as pq
import xxhash

parser = argparse.ArgumentParser()
parser.add_argument('--verify', action='store_true')
parser.add_argument('--output', type=Path, default=Path(__file__).parent)
args = parser.parse_args()
connection = duckdb.connect()
connection.execute('SET threads = 1')

with tempfile.TemporaryDirectory(prefix='columnar-bloom-') as work:
    parquet = Path(work) / 'bloom.parquet'
    connection.execute('''CREATE TABLE source AS
        SELECT i::BIGINT AS row_id, key::INTEGER AS key_i32,
            (9007199254740993 + key)::BIGINT AS key_i64,
            key::FLOAT AS key_f32, key::DOUBLE AS key_f64,
            'key-' || key::VARCHAR AS key_string,
            encode('key-' || key::VARCHAR) AS key_binary
        FROM (SELECT i, CASE WHEN i % 97 = 0 THEN NULL ELSE
            (i % 64) * 2 + (i // 2048) * 17 END AS key FROM range(6144) source(i))''')
    connection.execute('COPY source TO ? (FORMAT PARQUET, ROW_GROUP_SIZE 2048, COMPRESSION SNAPPY, BLOOM_FILTER_FALSE_POSITIVE_RATIO 0.25)', [str(parquet)])
    arrow = pq.read_table(parquet)
    connection.register('arrow_read', arrow)
    assert connection.execute('SELECT count(*) FROM ((FROM source EXCEPT ALL FROM arrow_read) UNION ALL (FROM arrow_read EXCEPT ALL FROM source))').fetchone()[0] == 0
    metadata = connection.execute('''SELECT row_group_id, path_in_schema, bloom_filter_offset, bloom_filter_length
        FROM parquet_metadata(?) ORDER BY row_group_id, column_id''', [str(parquet)]).fetchall()
    assert sorted(set(row[0] for row in metadata)) == [0, 1, 2]
    assert all(row[2] is not None for row in metadata if row[1] != 'row_id')

    def exclusions(column, value):
        return [row[2] for row in connection.execute('FROM parquet_bloom_probe(?, ?, ?) ORDER BY row_group_id', [str(parquet), column, value]).fetchall()]

    false_positives = [value for value in range(1, 126, 2) if not exclusions('key_i32', value)[0]]
    assert false_positives, 'The fixture must demonstrate an actual native false positive inside group bounds.'
    keys = [-1, 0, 1, 34, 42, 101, 160, false_positives[0]]
    probes = []
    for column in ['key_i32', 'key_i64', 'key_f32', 'key_f64', 'key_string', 'key_binary']:
        for key in dict.fromkeys(keys):
            value = 9007199254740993 + key if column == 'key_i64' else float(key) if column in ['key_f32', 'key_f64'] else f'key-{key}' if column == 'key_string' else f'key-{key}'.encode() if column == 'key_binary' else key
            ids = [row[0] for row in connection.execute(f'SELECT row_id FROM source WHERE {column} = ? ORDER BY row_id', [value]).fetchall()]
            serialized = value.hex() if isinstance(value, bytes) else str(value) if column == 'key_i64' else value
            probes.append({'column': column, 'value': serialized, 'excludes': exclusions(column, value), 'rowIds': ids})
    hashes = []
    for length in [0, 1, 3, 4, 7, 8, 15, 16, 31, 32, 33, 63, 64, 129]:
        value = bytes((index * 37 + 11) % 256 for index in range(length))
        hashes.append({'hex': value.hex(), 'xxh64': xxhash.xxh64(value, seed=0).hexdigest()})
    payload = parquet.read_bytes()
    proof = {'schemaVersion': 1, 'writer': 'DuckDB 1.4.4', 'oracle': 'DuckDB native Bloom probes / PyArrow 25.0.1 / xxhash 3.5.0',
             'file': {'name': 'bloom.parquet', 'size': len(payload), 'sha256': hashlib.sha256(payload).hexdigest()},
             'rows': arrow.num_rows, 'rowGroupRows': 2048, 'falsePositive': {'column': 'key_i32', 'value': false_positives[0], 'rowGroup': 0},
             'metadata': [{'rowGroup': row[0], 'column': row[1], 'offset': row[2], 'length': row[3]} for row in metadata],
             'probes': probes, 'hashes': hashes}
    outputs = {'bloom.parquet': payload, 'proof.json': (json.dumps(proof, indent=2) + '\n').encode()}
    args.output.mkdir(parents=True, exist_ok=True)
    for name, contents in outputs.items():
        path = args.output / name
        if args.verify:
            assert path.read_bytes() == contents, path
        else:
            with path.open('xb') as stream:
                stream.write(contents)
print(f'Verified {len(probes)} native probes, {len(hashes)} XXH64 cases, and false positive {false_positives[0]} in row group 0.')