"""
Scripts for postprocessing h5 files, e.g. shuffling.
"""
import os
import argparse
import warnings

import h5py
import km3pipe as kp
import km3modules as km
from orcasong.modules import EventSkipper
from orcasong.tools.concatenate import get_compopts, copy_attrs


def postproc_file(
        input_file,
        output_file=None,
        shuffle=True,
        event_skipper=None,
        delete=False,
        seed=42,
        statusbar_every=1000):
    """
    Postprocess a file using km3pipe after it has been preprocessed in OrcaSong.

    Parameters
    ----------
    input_file : str
        Path of the file that will be processed.
    output_file : str, optional
        If given, this will be the name of the output file.
        Otherwise, a name is auto generated.
    shuffle : bool
        Shuffle order of events.
    event_skipper : func, optional
        Function that takes the blob as an input, and returns a bool.
        If the bool is true, the event will be skipped.
    delete : bool
        Specifies if the input file should be deleted after processing.
    seed : int
        Sets a fixed random seed for the shuffling.
    statusbar_every : int or None
        After how many line a km3pipe status should be printed.

    Returns
    -------
    output_file : str
        Path to the output file.

    """
    if output_file is None:
        output_file = get_filepath_output(
            input_file, shuffle=shuffle, event_skipper=event_skipper)
    if os.path.exists(output_file):
        raise FileExistsError(output_file)

    print(f'Setting a Global Random State with the seed < {seed} >.')
    km.GlobalRandomState(seed=seed)

    comptopts = get_compopts(input_file)
    # km3pipe uses pytables for saving the shuffled output file,
    # which has the name 'zlib' for the 'gzip' filter
    if comptopts["complib"] == 'gzip':
        comptopts["complib"] = 'zlib'

    pipe = kp.Pipeline()
    if statusbar_every is not None:
        pipe.attach(km.common.StatusBar, every=statusbar_every)
        pipe.attach(km.common.MemoryObserver, every=statusbar_every)
    pipe.attach(
        kp.io.hdf5.HDF5Pump,
        filename=input_file,
        shuffle=shuffle,
        reset_index=True,
    )
    if event_skipper is not None:
        pipe.attach(EventSkipper, event_skipper=event_skipper)
    pipe.attach(
        kp.io.hdf5.HDF5Sink,
        filename=output_file,
        complib=comptopts["complib"],
        complevel=comptopts["complevel"],
        chunksize=comptopts["chunksize"],
        flush_frequency=1000,
    )
    pipe.drain()

    copy_used_files(input_file, output_file)
    copy_attrs(input_file, output_file)
    if delete:
        print("Deleting original file")
        os.remove(input_file)

    print("Done!")
    return output_file


def copy_used_files(source_file, target_file):
    """ Copy the "used_files" dataset from one h5 file to another, if it is present.
    """
    with h5py.File(source_file, "r") as src:
        if "used_files" in src:
            print("Copying used_files dataset")
            with h5py.File(target_file, "a") as trg:
                trg.create_dataset("used_files", data=src["used_files"])


def get_filepath_output(input_file, shuffle=True, event_skipper=None):
    """ Get the filename of the shuffled / rebalanced output file as a str.
    """
    fname_adtn = ''
    if shuffle:
        fname_adtn += '_shuffled'
    if event_skipper is not None:
        fname_adtn += '_reb'
    return f"{os.path.splitext(input_file)[0]}{fname_adtn}.h5"


def h5shuffle():
    # TODO deprecated
    warnings.warn("h5shuffle is deprecated and has been renamed to orcasong h5shuffle")
    parser = argparse.ArgumentParser(description='Shuffle an h5 file using km3pipe.')
    parser.add_argument('input_file', type=str, help='File to shuffle.')
    parser.add_argument('--output_file', type=str,
                        help='Name of output file. Default: Auto generate name.')
    parser.add_argument('--delete', action="store_true",
                        help='Delete original file afterwards.')

    postproc_file(**vars(parser.parse_args()), shuffle=True, event_skipper=None)
