#!/usr/bin/env python3
import sys
import os
import plistlib
import subprocess
from pprint import pprint


def disk_prewrite_darwin(dev_path):
    diskutil = subprocess.Popen([ 'sudo', 'diskutil', 'unmountDisk', dev_path ])
    diskutil.communicate()

    if diskutil.returncode != 0:
        print('Cannot unmount device {}'.format(dev_path))
        print('  diskutil returned error code {}'.format(diskutil.returncode))
        exit(-1)


def disk_prewrite(dev_path):
    if sys.platform == 'darwin':
        disk_prewrite_darwin(dev_path)


def disk_postwrite_darwin(dev_path):
    sync = subprocess.Popen([ 'sync' ])
    sync.communicate()

    diskutil = subprocess.Popen([ 'sudo', 'diskutil', 'unmountDisk', dev_path ])
    diskutil.communicate()


def disk_postwrite_linux(dev_path):
    sync = subprocess.Popen([ 'sync' ])
    sync.communicate()


def disk_postwrite(dev_path):
    if sys.platform == 'darwin':
        disk_postwrite_darwin(dev_path);
    elif sys.platform == 'linux':
        disk_postwrite_linux(dev_path)


def get_devices_darwin():
    diskutil = subprocess.Popen([ 'diskutil', 'list', '-plist' ], stdout=subprocess.PIPE)
    (out, err) = diskutil.communicate()

    if diskutil.returncode != 0:
        print('Cannot query storage devices')
        print('  diskutil returned error code {}'.format(diskutil.returncode))
        print(err)
        exit(-1)

    all_device_data = plistlib.loads(out)
    device_names = []

    for v in all_device_data['AllDisksAndPartitions']:
        if v['OSInternal']:
            continue

        device_names.append(v['DeviceIdentifier'])

    final_list = []

    for dev in device_names:
        diskutil = subprocess.Popen([ 'diskutil', 'info', '-plist', dev ], stdout=subprocess.PIPE)
        (out, err) = diskutil.communicate()

        if diskutil.returncode != 0:
            print('Cannot query device {}'.format(dev))
            print('  diskutil returned error code {}'.format(diskutil.returncode))
            print(err)
            exit(-1)

        dev_info = plistlib.loads(out)

        if dev_info['BusProtocol'] != 'USB':
            continue

        final_list.append(('/dev/r{}'.format(dev_info['DeviceIdentifier']), dev_info['MediaName']))

    return final_list


def get_devices_linux():
    final_list = []
    dev_list = os.listdir('/sys/block')

    for dev_name in dev_list:
        with open('/sys/block/{}/removable'.format(dev_name), 'r') as dev_removable:
            is_removable = int(dev_removable.read(1))
            if is_removable != 1:
                continue

        with open('/sys/block/{}/device/model'.format(dev_name)) as dev_model_name:
            model_name = dev_model_name.read().strip()
            final_list.append(('/dev/{}'.format(dev_name), model_name))

    return final_list


def get_devices():
    if sys.platform == 'darwin':
        return get_devices_darwin()
    elif sys.platform == 'linux':
        return get_devices_linux()
    else:
        print('This script cannot be used on {}'.format(sys.platform))
        exit(-1)


def choice_options(num_devices):
    if num_devices == 1:
        return '1'

    return '1-{}'.format(num_devices)


if not os.path.isfile('build/socks-kernel.iso'):
    print('No system ISO image found.')
    print('Please run \'make cd\' to generate an ISO image')
    exit(-1)

devices = get_devices()

if len(devices) == 0:
    print('No devices are available for installation')
    print('Note: this script only considers removable USB drives for installation')
    exit(-1)

print('Available Devices:')
for i, dev in enumerate(devices):
    print(' {}. {} ({})'.format(i + 1, dev[1], dev[0]))

choice = 0
while True:
    print('Selection [{}]: '.format(choice_options(len(devices))), end='')
    choice_str = input()

    try:
        choice = int(choice_str)
    except:
        print(' {} is not a valid option'.format(choice_str))
        continue

    if choice < 1 or choice > len(devices):
        print(' {} is not a valid option'.format(choice_str))
        continue

    choice -= 1
    break

print('Installing to {} at {}'.format(devices[choice][1], devices[choice][0]))

disk_prewrite(devices[choice][0])

dd_args = [
    'sudo',
    'dd',
    'if=build/socks-kernel.iso',
    'of={}'.format(devices[choice][0]),
    'bs=1{}'.format('m' if sys.platform == 'darwin' else 'M')
]

dd = subprocess.Popen(dd_args)
dd.communicate()
if dd.returncode != 0:
    print('Cannot write ISO file to device {}'.format(devices[choice][0]))
    print('  dd returned error code {}'.format(dd.returncode))

disk_postwrite(devices[choice][0])