import json
import cv2
import numpy as np

class aif:
    def __init__(self, dpath, ipath):
        self.data_fh = None
        self.data_idx = None
        self.datainfo = None
        self.cursor = 0
        self.dpath = dpath
        self.ipath = ipath

    def wbegin(self):
        self.data_fh = open(self.dpath, 'wb')
        self.idx_fh = open(self.ipath, 'w')
        self.datainfo = []
        self.cursor = 0

    def rbegin(self):
        self.data_fh = open(self.dpath, 'rb')
        self.idx_fh = open(self.ipath, 'r')
        self.datainfo = json.load(self.idx_fh)
        return self.datainfo

    def append(self, img):
        try:
            success, enc_img = cv2.imencode(".png", img)
            data = enc_img.tobytes()
            self.data_fh.write(data)
            self.datainfo.append([self.cursor, len(data)])
            self.cursor += len(data)
        except Exception as err:
            print("libAIF append() exception: " + str(err))
            self.datainfo.append([self.cursor, 0])

    def append_any(self, data):
        try:
            self.data_fh.write(data)
            self.datainfo.append([self.cursor, len(data)])
            self.cursor += len(data)
        except Exception as err:
            print("libAIF append() exception: " + str(err))
            self.datainfo.append([self.cursor, 0])

    def extract(self, idx):
        xinfo = self.datainfo[idx]
        self.data_fh.seek(xinfo[0])
        img_dat = self.data_fh.read(xinfo[1])
        img_dat = np.asarray(bytearray(img_dat), dtype='uint8')
        if len(img_dat) > 3:
            img = cv2.imdecode(img_dat, cv2.IMREAD_COLOR)
        else:
            return np.array([])
        return img
    
    def extract_raw(self, idx, dtype='uint8'):
        xinfo = self.datainfo[idx]
        self.data_fh.seek(xinfo[0])
        img_dat = self.data_fh.read(xinfo[1])
        # img_dat = np.asarray(bytearray(img_dat), dtype='uint8')
        img_dat = np.frombuffer(img_dat, dtype=dtype)
        return img_dat

    def __finish(self):
        self.data_fh.close()
        self.idx_fh.close()

    def wfinish(self):
        json.dump(self.datainfo, self.idx_fh)
        self.__finish()

    def rfinish(self):
        self.__finish()

if __name__ == "__main__":
    xaif = aif("D:/ecg/s2/f_imgs1/datafile", "D:/ecg/s2/f_imgs1/idxfile")
    datinfo = xaif.rbegin()
    for i in range(0, len(datinfo)):
        cv2.imwrite("D:/ecg/s2/f_imgs1/" + str(i) + ".png", xaif.extract(i))
    xaif.rfinish()