import sys
import os
import argparse
import nibabel as nb


def lossless_slice(img, slicers):
    if not nb.imageclasses.spatial_axes_first(img):
        raise ValueError("Cannot slice an image that is not known to have spatial axes first")

    scaling = hasattr(img.header, 'set_slope_inter')

    data = img.dataobj._get_unscaled(slicers) if scaling else img.dataobj[slicers]
    roi_img = img.__class__(data, affine=img.slicer.slice_affine(slicers), header=img.header)

    if scaling:
        roi_img.header.set_slope_inter(img.dataobj.slope, img.dataobj.inter)
    return roi_img


def parse_slice(crop, allow_step=True):
    if crop is None:
        return slice(None)
    start, stop, *extra = [int(val) if val else None for val in crop.split(":")]
    if len(extra) > 1:
        raise ValueError(f"Cannot parse specification: {crop}")
    if not allow_step and extra and extra[0] not in (1, None):
        raise ValueError(f"Step entry not permitted: {crop}")

    step = extra[0] if extra else None
    if step not in (1, -1, None):
        raise ValueError(f"Downsampling is not supported: {crop}")

    return slice(start, stop, step)


def sanitize(args):
    # Argparse likes to treat "-1:..." as a flag
    return [f' {arg}' if arg[0] == '-' and ":" in arg else arg
            for arg in args]


def main(args=None):
    if args is None:
        args = sys.argv[1:]
    parser = argparse.ArgumentParser(
        description="Crop images to a region of interest",
        epilog="If a start or stop value is omitted, the start or end of the axis is assumed.")
    parser.add_argument('--version', action='version', version=nb.__version__)
    parser.add_argument("-i", metavar="I1:I2[:-1]",
                        help="Start/stop [flip] along first axis (0-indexed)")
    parser.add_argument("-j", metavar="J1:J2[:-1]",
                        help="Start/stop [flip] along second axis (0-indexed)")
    parser.add_argument("-k", metavar="K1:K2[:-1]",
                        help="Start/stop [flip] along third axis (0-indexed)")
    parser.add_argument("-t", metavar="T1:T2", help="Start/stop along fourth axis (0-indexed)")
    parser.add_argument("in_file", help="Image file to crop")
    parser.add_argument("out_file", help="Output file name")

    opts = parser.parse_args(args=sanitize(args))

    try:
        islice = parse_slice(opts.i)
        jslice = parse_slice(opts.j)
        kslice = parse_slice(opts.k)
        tslice = parse_slice(opts.t, allow_step=False)
    except ValueError as err:
        print(f"Could not parse input arguments. Reason follows.\n{err}")
        return 1

    kwargs = {}
    if os.path.realpath(opts.in_file) == os.path.realpath(opts.out_file):
        kwargs['mmap'] = False
    img = nb.load(opts.in_file, **kwargs)

    slicers = (islice, jslice, kslice, tslice)[:img.ndim]
    expected_shape = nb.fileslice.predict_shape(slicers, img.shape)
    if any(dim == 0 for dim in expected_shape):
        print(f"Cannot take zero-length slices. Predicted shape {expected_shape}.")
        return 1

    try:
        sliced_img = lossless_slice(img, slicers)
    except Exception:
        print("Could not slice image. Full traceback follows.")
        raise
    nb.save(sliced_img, opts.out_file)
    return 0
