#!/usr/bin/env python3
"""Minimal pure-Python Parquet reader (flat schemas, PLAIN/dictionary pages,
SNAPPY/GZIP/uncompressed). Enough to extract tick columns from public datasets.
"""
import struct
import sys
import zlib
from array import array

# ---------------------------------------------------------------- thrift compact
class TReader:
    def __init__(self, buf, pos=0):
        self.b = buf
        self.p = pos

    def byte(self):
        v = self.b[self.p]
        self.p += 1
        return v

    def varint(self):
        shift = 0
        result = 0
        while True:
            c = self.b[self.p]
            self.p += 1
            result |= (c & 0x7F) << shift
            if not (c & 0x80):
                return result
            shift += 7

    def zigzag(self):
        n = self.varint()
        return (n >> 1) ^ -(n & 1)

    def read_value(self, t):
        if t == 1:
            return True
        if t == 2:
            return False
        if t == 3:
            v = self.byte()
            return v - 256 if v > 127 else v
        if t in (4, 5, 6):
            return self.zigzag()
        if t == 7:
            v = struct.unpack_from('<d', self.b, self.p)[0]
            self.p += 8
            return v
        if t == 8:
            n = self.varint()
            v = bytes(self.b[self.p:self.p + n])
            self.p += n
            return v
        if t in (9, 10):
            h = self.byte()
            size = h >> 4
            et = h & 15
            if size == 15:
                size = self.varint()
            out = []
            for _ in range(size):
                if et in (1, 2):
                    out.append(self.byte() == 1)
                else:
                    out.append(self.read_value(et))
            return out
        if t == 11:
            size = self.varint()
            out = {}
            if size:
                kv = self.byte()
                for _ in range(size):
                    k = self.read_value(kv >> 4)
                    out[k] = self.read_value(kv & 15)
            return out
        if t == 12:
            return self.read_struct()
        raise ValueError('bad thrift type %d' % t)

    def read_struct(self):
        out = {}
        fid = 0
        while True:
            h = self.byte()
            if h == 0:
                return out
            t = h & 15
            d = h >> 4
            if d == 0:
                fid = self.zigzag()
            else:
                fid += d
            out[fid] = self.read_value(t)


# ---------------------------------------------------------------- snappy
def snappy_decompress(src):
    r = TReader(src)
    n = r.varint()
    out = bytearray()
    p = r.p
    L = len(src)
    while p < L:
        tag = src[p]
        p += 1
        kind = tag & 3
        if kind == 0:
            ln = tag >> 2
            if ln >= 60:
                nb = ln - 59
                ln = int.from_bytes(src[p:p + nb], 'little')
                p += nb
            ln += 1
            out += src[p:p + ln]
            p += ln
            continue
        if kind == 1:
            ln = ((tag >> 2) & 7) + 4
            off = ((tag >> 5) << 8) | src[p]
            p += 1
        elif kind == 2:
            ln = (tag >> 2) + 1
            off = src[p] | (src[p + 1] << 8)
            p += 2
        else:
            ln = (tag >> 2) + 1
            off = int.from_bytes(src[p:p + 4], 'little')
            p += 4
        start = len(out) - off
        if off >= ln:
            out += out[start:start + ln]
        else:
            while ln > 0:
                chunk = out[start:start + min(off, ln)]
                out += chunk
                ln -= len(chunk)
    if len(out) != n:
        raise ValueError('snappy length mismatch')
    return bytes(out)


_zstd = None


def zstd_decompress(data, usize):
    """ZSTD pages are decoded with the system libzstd (or $LIBZSTD) via ctypes."""
    global _zstd
    import ctypes
    import ctypes.util
    import os
    if _zstd is None:
        lib = os.environ.get('LIBZSTD') or ctypes.util.find_library('zstd')
        if not lib:
            raise RuntimeError('libzstd not found: install it or set LIBZSTD=/path/to/libzstd.so')
        _zstd = ctypes.CDLL(lib)
        _zstd.ZSTD_decompress.restype = ctypes.c_size_t
        _zstd.ZSTD_decompress.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t]
        _zstd.ZSTD_isError.restype = ctypes.c_uint
        _zstd.ZSTD_isError.argtypes = [ctypes.c_size_t]
    dst = ctypes.create_string_buffer(usize)
    src = bytes(data)
    r = _zstd.ZSTD_decompress(dst, usize, src, len(src))
    if _zstd.ZSTD_isError(r) or r != usize:
        raise ValueError('zstd failure')
    return dst.raw


def decompress(codec, data, usize):
    if codec == 6:
        return zstd_decompress(data, usize)
    if codec == 0:
        return data
    if codec == 1:
        return snappy_decompress(data)
    if codec == 2:
        return zlib.decompress(data, 16 + zlib.MAX_WBITS)
    raise ValueError('unsupported codec %d' % codec)


# ---------------------------------------------------------------- rle hybrid
def rle_hybrid(buf, pos, end, width, count):
    out = []
    r = TReader(buf, pos)
    mask = (1 << width) - 1
    vbytes = (width + 7) // 8
    while len(out) < count and r.p < end:
        h = r.varint()
        if h & 1:
            groups = h >> 1
            nb = groups * width
            chunk = int.from_bytes(buf[r.p:r.p + nb], 'little')
            r.p += nb
            if width == 0:
                out.extend([0] * (groups * 8))
            else:
                for i in range(groups * 8):
                    out.append((chunk >> (i * width)) & mask)
        else:
            run = h >> 1
            v = int.from_bytes(buf[r.p:r.p + vbytes], 'little') if vbytes else 0
            r.p += vbytes
            out.extend([v] * run)
    return out[:count]


PTYPE_FMT = {1: ('i', 4), 2: ('q', 8), 4: ('f', 4), 5: ('d', 8)}


def plain_values(ptype, buf, pos, count):
    code, size = PTYPE_FMT[ptype]
    a = array(code)
    a.frombytes(bytes(buf[pos:pos + count * size]))
    if sys.byteorder != 'little':
        a.byteswap()
    return a


def read_parquet(path, wanted=None):
    with open(path, 'rb') as f:
        data = f.read()
    if data[:4] != b'PAR1' or data[-4:] != b'PAR1':
        raise ValueError('not parquet')
    mlen = struct.unpack_from('<I', data, len(data) - 8)[0]
    meta = TReader(data, len(data) - 8 - mlen).read_struct()
    schema = meta[2]
    leaves = []
    for el in schema[1:]:
        if el.get(5):
            raise ValueError('nested schema not supported')
        leaves.append((el[4].decode(), el.get(1), el.get(3, 0), el.get(6), el.get(10)))
    cols = {name: [] for name, *_ in leaves}
    for rg in meta[4]:
        for ci, cc in enumerate(rg[1]):
            name, ptype, rep, conv, logical = leaves[ci]
            if wanted and name not in wanted:
                continue
            md = cc[3]
            codec = md[4]
            nvals = md[5]
            start = md.get(11) or md[9]
            if md.get(11) is not None and md[11] < md[9]:
                start = md[11]
            pos = start
            dictionary = None
            got = 0
            while got < nvals:
                r = TReader(data, pos)
                ph = r.read_struct()
                pos = r.p
                ptype_page = ph[1]
                usize = ph[2]
                csize = ph[3]
                raw = data[pos:pos + csize]
                pos += csize
                if ptype_page == 2:  # dictionary
                    page = decompress(codec, raw, usize)
                    dictionary = plain_values(ptype, page, 0, ph[7][1])
                    continue
                if ptype_page == 0:
                    page = decompress(codec, raw, usize)
                    dh = ph[5]
                    n = dh[1]
                    enc = dh[2]
                    p = 0
                    ndef = n
                    if rep == 1:  # optional -> def levels
                        ln = struct.unpack_from('<I', page, 0)[0]
                        lv = rle_hybrid(page, 4, 4 + ln, 1, n)
                        p = 4 + ln
                        ndef = sum(lv)
                        if ndef != n:
                            raise ValueError('nulls not supported')
                elif ptype_page == 3:
                    dh = ph[8]
                    n = dh[1]
                    if dh.get(2, 0):
                        raise ValueError('nulls not supported')
                    enc = dh[4]
                    dl = dh.get(5, 0)
                    rl = dh.get(6, 0)
                    lev = raw[:dl + rl]
                    body = raw[dl + rl:]
                    if dh.get(7, True):
                        body = decompress(codec, body, usize - dl - rl)
                    page = body
                    p = 0
                    ndef = n
                else:
                    continue
                if enc == 0:
                    vals = plain_values(ptype, page, p, ndef)
                elif enc in (2, 8):
                    width = page[p]
                    idx = rle_hybrid(page, p + 1, len(page), width, ndef)
                    d = dictionary
                    vals = [d[i] for i in idx]
                else:
                    raise ValueError('unsupported encoding %d' % enc)
                cols[name].extend(vals)
                got += n
    return cols, leaves, meta[3]


if __name__ == '__main__':
    cols, leaves, nrows = read_parquet(sys.argv[1])
    print('rows', nrows)
    for name, ptype, rep, conv, logical in leaves:
        c = cols[name]
        print(name, ptype, rep, conv, logical, len(c), list(c[:5]))