import sys
import zlib
import json
import base64
import requests
import tempfile
import contextlib
import pyarrow as pa
import OpenSSL.crypto


@contextlib.contextmanager
def pfx_to_pem(pfx_path, pfx_password):
    '''Decrypts the .pfx file to be used with requests.'''
    with tempfile.NamedTemporaryFile(suffix='.pem') as t_pem:
        t_pem.close()
        f_pem = open(t_pem.name, 'wb')
        pfx = open(pfx_path, 'rb').read()
        p12 = OpenSSL.crypto.load_pkcs12(pfx, pfx_password)
        f_pem.write(OpenSSL.crypto.dump_privatekey(OpenSSL.crypto.FILETYPE_PEM, p12.get_privatekey()))
        f_pem.write(OpenSSL.crypto.dump_certificate(OpenSSL.crypto.FILETYPE_PEM, p12.get_certificate()))
        ca = p12.get_ca_certificates()
        if ca is not None:
            for cert in ca:
                f_pem.write(OpenSSL.crypto.dump_certificate(OpenSSL.crypto.FILETYPE_PEM, cert))
        f_pem.close()
        yield t_pem.name


base_url = '<your-base-url>'  # from Liberator Profile page download
compressed_transfer = True

# POST the query and prepare for a stream of single-line JSON replies
with pfx_to_pem('liberator.pfx', '') as cert:
    r = requests.post(f'{base_url}/liberator/query',
            data=json.dumps({"compress": compressed_transfer, "json_xfer": False,
                "user": user, "token": token,
                "name": name, "as_of": as_of, "back_to": back_to, "symbols": symbols,
                "system": "API"}),
            headers={'Content-Type': 'application/json'},
            cert=cert,
            stream=True)

# Ensure that the request was successful
if r.status_code != 200:
    print(r.json())
else:
    batches = []
    for line in r.iter_lines(4096):
        # Show progress
        print('.', end='')
        sys.stdout.flush()
        # The response stream contains informational messages other than "batch"
        if line[0:14] == b'{"exception":"':
            print('Error: ' + json.loads(''.join([chr(c) for c in line]), encoding='latin-1')['value'])
            break
        if line[0:10] != b'{"batch":"':
            continue
        # Cut out the payload from the JSON line without json module overhead
        decoded = base64.b64decode(line[10:-3])
        # Interpret the payload as an Apache Arrow IPC stream
        reader = pa.ipc.open_stream(zlib.decompress(decoded) if compressed_transfer else decoded)
        # Accumulate RecordBatch objects in a list
        batches.extend([batch for batch in reader])
    print('')
    if batches:
        # Create an Arrow Table view on the batches and convert to a pandas DataFrame
        df = pa.Table.from_batches(batches).to_pandas()
        if 'muts' in df and '_seq' in df:
            # By default, the data is sorted by symbol... re-sort by time
            df = df.sort_values(['muts', '_seq'])
            df.reset_index(drop=True, inplace=True)
        print(df)