#!/usr/bin/env fslpython


import bz2
import shlex
import string
import os.path as op
import itertools as it
import subprocess as sp

from io import BytesIO

import numpy as np

from fsl.utils.tempdir import tempdir


FORMATS = ['npy', 'gz', 'zst', 'bz2']


def run(cmd, **kwargs):
    print(f'RUN: {cmd}')

    kwargs['check'] = kwargs.get('check', True)

    if 'stdout' not in kwargs:
        kwargs['capture_output'] = True

    try:
        return sp.run(shlex.split(cmd), **kwargs)
    except Exception:
        print('Error running command! If testing locally, make '
              'sure the fsl_testznz binary is on your $PATH')
        raise


def load_test_file(filename):

    fmt = op.splitext(filename)[1].strip('.')

    with open(filename, 'rb') as f:
        data = f.read()

    if fmt == 'gz':
        data = run(f'gzip  -dc', input=data).stdout
    elif fmt == 'zst':
        data = run(f'zstd  -dc', input=data).stdout
    elif fmt == 'bz2':
        data = run(f'bzip2 -dc', input=data).stdout

    print(f'Load {filename} [{len(data)} bytes]')

    return data

def check_file(filename):
    """Use the underlying compression command to check file integrity."""
    fmt = op.splitext(filename)[1].strip('.')

    if   fmt == 'gz':  cmd = f'gzip  --test {filename}'
    elif fmt == 'zst': cmd = f'zstd  --test {filename}'
    elif fmt == 'bz2': cmd = f'bzip2 --test {filename}'
    else: raise Exception(f'unknown format {fmt}')

    result = run(cmd, check=False)

    if result.returncode != 0:
        raise Exception(f'{result.stdout} - {result.stderr}')


def create_test_file(filename, data):

    fmt = op.splitext(filename)[1].strip('.')

    print(f'Create {filename} [{len(data)} bytes]')

    with open(filename, 'wb') as f:
        if fmt == 'gz':
            run('gzip  -c', input=data, stdout=f)
        elif fmt == 'zst':
            run('zstd  -c', input=data, stdout=f)
        elif fmt == 'bz2':
            run('bzip2 -c', input=data, stdout=f)
        else:
            f.write(data)


def test_znz_guessformat():

    data = (string.ascii_letters + string.digits).encode()

    with tempdir():
        for fmt in FORMATS:
            infile = f'input.{fmt}'
            create_test_file(infile, data)
            got = run(f'fsl_testznz -i {infile}').stdout
            assert got == data


def test_znz_roundtrip():

    # square matrix approximately 50MiB uncompressed
    dim  = int(np.ceil(np.sqrt((50 * 1048576) // 8)))
    data = np.random.random((dim, dim)).astype(np.float64)

    encoded = BytesIO()
    np.save(encoded, data, allow_pickle=False)
    encoded = encoded.getvalue()

    with tempdir():

        for infmt, outfmt in it.product(FORMATS, FORMATS):

            infile  = f'input.{infmt}'
            outfile = f'output.{outfmt}'

            create_test_file(infile, encoded)

            run(f'fsl_testznz -i {infile} -o {outfile}')

            got = load_test_file(outfile)
            got = np.load(BytesIO(got))

            assert np.all(data == got)


def test_znz_output_file_integrity():

    # fsl/znzlib!5
    #
    # zstd write buffer size defaults to 131072 bytes
    # (see ZSTD_CStreamInSize)
    data = np.random.random((160, 224, 256)).astype(np.float32)

    with tempdir():

        data.tofile('input.raw')
        for fmt in ['gz', 'zst', 'bz2']:
            run(f'fsl_testznz -i input.raw -o output.{fmt}')
            check_file(f'output.{fmt}')


def test_znz_seek():

    header = b'1234567890'
    footer = (string.ascii_letters + string.digits).encode()
    data   = b'1' * 1048576

    with tempdir():

        for fmt in FORMATS:
            infile = f'input.{fmt}'
            create_test_file(infile, header + data + footer)
            got = run(f'fsl_testznz -i {infile} -s "0,10;1048586,-1"').stdout
            assert got == header + footer


def test_znz_seek_relative():

    header  = b'1234567890'
    footer  = (string.ascii_letters + string.digits).encode()
    padding = b'0' * 1048576
    data    = b'1' * 25

    with tempdir():

        for fmt in FORMATS:
            infile = f'input.{fmt}'
            create_test_file(infile, header + padding + data + padding + footer)
            got = run(f'fsl_testznz -i {infile} -r -s 0,10;1048576,25;1048576,-1').stdout
            assert got == header + data + footer


if __name__ == '__main__':
    test_znz_guessformat()
    test_znz_roundtrip()
    test_znz_seek()
    test_znz_seek_relative()
    test_znz_output_file_integrity()
