import subprocess

import os

SINGLE_DISK_OFFSET = 1572864

MAX_IMAGES_ON_DRIVE = 1000

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:

        raise Exception('Not created')

    return filename

def format_images(device, indexes, temp_path, size=1440):

    for ipair in indexes:

        if isinstance(ipair, int):

            image = get_or_create_empty_image(ipair, temp_path, size)

            write_image(device, image, 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)

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[0x2B:0x36].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