import struct
import argparse
import sys
import os

from shlex import split as shlex_split
from subprocess import Popen, PIPE

sg_raw = os.path.join(os.path.dirname(__file__), 'sg_raw.exe')
sg_verify = os.path.join(os.path.dirname(__file__), 'sg_verify.exe')

PACKET_SIZE = 16384

def upload(path, device_id):
    if not os.path.exists(path) or os.path.isdir(path):
        print(f'\'{path}\' not found')
        return

    buf = ''
    with open(path, 'rb') as f:
        buf = f.read()
        f.close()

    buf_len = len(buf)

    packet_leftovers = buf_len % PACKET_SIZE
    if packet_leftovers > 0:
        print('Padding buffer')
        padding_size = PACKET_SIZE - packet_leftovers
        buf += struct.pack('%ss' % padding_size, '\xFF'*padding_size)

    buf_len = len(buf)
    packets = buf_len // PACKET_SIZE
    print(f'Sending {packets} packets to {device_id}')

    send_header = '00'
    last_total = 0

    for i in range(0, packets):
        cmd_str = f'{sg_raw} {device_id} -v -b -i {path} -k {last_total} -s {PACKET_SIZE} -n 2A 00 00 00 00 {send_header} 00 00 20 00'
        cmd_str = cmd_str.replace('\\', '/')
        cmd = shlex_split(cmd_str)
        send_header = '01'
        p = Popen(cmd, stdin=PIPE, stderr=PIPE)
        err = p.stderr.read().decode("utf-8")

        percent = int(last_total * 100 / buf_len)
        print(f'\r[{percent}%]', end='', flush=True)

        if not 'Good' in err:
            print(f'SCSI error: {err}')
            return
        
        last_total += PACKET_SIZE

    print('\r[100%]')
    p = Popen([sg_verify, device_id], stderr=PIPE)
    err = p.stderr.read()

    if len(err) != 0:
        print(f'SCSI error: {err}')
        return

if __name__ == '__main__':
    parser = argparse.ArgumentParser('upload')
    parser.add_argument('-f', '--file')
    parser.add_argument('-d', '--device')
    args = parser.parse_args()

    if args.file is None:
        print('no file specified')
        sys.exit()
    if args.device is None:
        print('no device specified')
        sys.exit()

    print(f'found sg_raw at {sg_raw}')
    print(f'found sg_verify at {sg_verify}')
    upload(args.file, args.device)
    print("success!")
