import subprocess
import os


SINGLE_DISK_OFFSET = 1572864
MAX_IMAGES_ON_DRIVE = 1000
LABEL_START_OFFSET = 0x2B
LABEL_LENGTH = 11


def get_or_create_empty_image(index, temp_path, size=1440):
    filename = 'empty{}.ima'.format(index)
    filename = os.path.join(temp_path, filename)

    if os.path.isfile(filename):
        os.remove(filename)
    proc = subprocess.Popen(['mkfs.msdos', '-C', filename, str(size)])
    rval = proc.wait()
    if rval != 0:
        print(proc.stdout.read())
        print(proc.stderr.read())
        raise Exception('Not created')

    return filename


def format_images(device, indexes, temp_path, size=1440, callback=None):
    for ipair in indexes:
        if isinstance(ipair, int):
            image = get_or_create_empty_image(ipair, temp_path, size)
            write_image(device, image, ipair)
            os.remove(image)
            if callback:
                callback(indexes, ipair)
        else:
            for i in range(ipair[0], ipair[1]):
                image = get_or_create_empty_image(i, temp_path, size)
                write_image(device, image, i)
                os.remove(image)
                if callback:
                    callback(ipair, i)


def write_image(device, path, index=0):
    """
    Writes image from path to device at specified index
    :param device:
    :param index:
    :param path:
    :return:
    """
    with open(device, 'r+b') as dfile:
        offset = SINGLE_DISK_OFFSET * index
        dfile.seek(offset)
        with open(path, 'rb') as ifile:
            fdata = ifile.read()
            dfile.write(fdata)


def read_image(device, path, index=0):
    """
    Reads image at specified index from device into path
    :param device:
    :param path:
    :param index:
    :return:
    """
    with open(device, 'rb') as dfile:
        offset = SINGLE_DISK_OFFSET * index
        dfile.seek(offset)
        bootsector = dfile.read(512)
        if bootsector[-1] == 0xAA and bootsector[-2] == 0x55:
            # FAT12 image
            sectorsize = bootsector[0xB] + (bootsector[0xC] << 8)
            sectorcount = bootsector[0x13] + (bootsector[0x14] << 8)
            with open(path, 'xb') as ifile:
                dfile.seek(offset)
                fdata = dfile.read(sectorsize * sectorcount)
                ifile.write(fdata)


def list_images(device):
    """
    Reads device and lists all of images on it
    :param device:
    :return: (label, size)
    """
    rval = []
    with open(device, 'rb') as dfile:
        dfile.seek(0, 2)
        disks = min(dfile.tell() / SINGLE_DISK_OFFSET, MAX_IMAGES_ON_DRIVE)
        dfile.seek(0)
        while len(rval) < disks:
            bootsector = dfile.read(512)
            label = None
            if bootsector[0] == 0xEB or bootsector[0] == 0xE9:
                sectorsize = bootsector[0xB] + (bootsector[0xC] << 8)
                sectorcount = bootsector[0x13] + (bootsector[0x14] << 8)
                disksize = sectorsize * sectorcount
                if bootsector[0x26] == 0x29:
                    # label probably exists.
                    # As an additional precaution check fs type.
                    # If it's valid, then get label, otherwise consider disk unlabeled
                    if bootsector[0x36:0x3E] in (b'FAT12   ', b'FAT16   ', b'FAT     '):
                        label = bootsector[LABEL_START_OFFSET:LABEL_START_OFFSET + LABEL_LENGTH].decode(
                            encoding='cp1251'
                        )
                dfile.seek(SINGLE_DISK_OFFSET - 512, 1)
            else:
                disksize = 0
                dfile.seek(SINGLE_DISK_OFFSET - 512, 1)
            rval.append((label, disksize))
    return rval


def update_label(device, index, label):
    offset = SINGLE_DISK_OFFSET * index
    label = label[0:11]
    label = label + ' ' * (LABEL_LENGTH - len(label))
    with open(device, 'wb') as dfile:
        dfile.seek(offset + LABEL_START_OFFSET, 0)
        dfile.write(label.encode(encoding='cp1251'))