import lib.fish_eye.pyARIS as pyARIS
import struct

def crop_clip(aris_path, num_frames, verbose=False):
    """
    Crop an aris file based on the first *num_frames* frames. Save the new aris files in the tmp folder
    """

    # load aris file and extract frame size
    ARIS_data, frame = pyARIS.DataImport(aris_path)
    FrameSize = ARIS_data.NumRawBeams*ARIS_data.SamplesPerChannel
    if verbose: print("True Old", ARIS_data.FrameCount, ARIS_data.StartFrame, ARIS_data.EndFrame)

    # get byte index of the cutoff point
    frameoffset = (1024+(num_frames*(1024+(FrameSize))))

    # read aris the bytes for the head and frames we want and cast to bytearray
    data = open(ARIS_data.filename, 'rb')
    cropped = data.read(frameoffset)
    array = bytearray(cropped)

    # get old values for important metadata
    old_frame_count     = struct.unpack("I", array[4:8])[0]
    old_start_frame     = struct.unpack("I", array[352:356])[0]
    old_end_frame       = struct.unpack("I", array[356:360])[0]

    if verbose: print("old", old_frame_count, old_start_frame, old_end_frame)

    # set new values
    array[4:8]        = bytearray(struct.pack("I", num_frames))
    array[352:356]    = bytearray(struct.pack("I", old_start_frame))
    array[356:360]    = bytearray(struct.pack("I", old_start_frame + num_frames))
    if verbose: print("new", array[4:8], array[352:356], array[356:360])
    
    # cast to bytes
    cropped = bytes(array)

    # save new aris file
    with open("tmp/cropped_aris.aris", 'wb') as f:
        f.write(cropped)

    # load file to check that the frame count, start frame and end frame makes sense
    if verbose: 
        ARIS_data_2, frame = pyARIS.DataImport("tmp/cropped_aris.aris")
        print("check", ARIS_data_2.FrameCount, ARIS_data_2.StartFrame, ARIS_data_2.EndFrame)