import getpass
import glob
import platform
from io import BytesIO
from os import path

import psutil

system_name = platform.system()


class DeviceNotFound(Exception):
    pass


def fetch():
    """获取U盘中的Excel文件"""
    if system_name == 'Windows':
        usbs = [disk.device for disk in psutil.disk_partitions() if disk.opts == 'rw,removable']
        if not usbs:
            raise DeviceNotFound()
        usb = usbs[0]
        files = glob.glob(f'{usb}/*.xlsx') + glob.glob(f'{usb}/*.xls')
        return [{"name": path.basename(f), "file": f} for f in files]

    usbs = [item for item in psutil.disk_partitions() if item.mountpoint.startswith(f'/media/{getpass.getuser()}')]
    if not usbs:
        raise DeviceNotFound()
    usb = usbs[0]
    files = glob.glob(f'{usb.mountpoint}/*.xlsx') + glob.glob(f'{usb.mountpoint}/*.xls')
    return [{"name": path.basename(f), "file": f} for f in files]


def put_in(filename: str, bin: BytesIO):
    if system_name == 'Windows':
        usbs = [disk.device for disk in psutil.disk_partitions() if disk.opts == 'rw,removable']
        if not usbs:
            raise DeviceNotFound()
        path_file = path.join(usbs[0], filename)
    else:
        usbs = [item for item in psutil.disk_partitions() if item.mountpoint.startswith(f'/media/{getpass.getuser()}')]
        if not usbs:
            raise DeviceNotFound()
        path_file = path.join(usbs[0].mountpoint, filename)
    with open(path_file, 'wb') as f:
        f.write(bin.getbuffer())