import threading
import os

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

erasing_methods = {
    "2 прохода": {
        "num_passes": 2,
        "data": [b"\x00", "random"]
    },
    "Случайные данные": {
        "num_passes": 1,
        "data": ["random"]
    },
    "Гутмана": {
        "num_passes": 0,
        "data": []
    },
    "3 прохода": {
        "num_passes": 3,
        "data": [b"\x00", b"\xff", "random"]
    },
    "7 проходов": {
        "num_passes": 7,
        "data": [b"\x00", b"\xff", "random", "random", b"\x00", b"\x01", "random"]
    }
}

def get_random_bytes(size):
    seed=os.urandom(32)
    backend = default_backend()
    cipher = Cipher(algorithms.AES(seed), modes.CTR(b'\0'*16), backend=backend)
    encryptor = cipher.encryptor()

    nulls=b'\0'*(size)
    return encryptor.update(nulls)

class Drive():

    index: int = None

    path: str

    total_sectors: str

    disk_type = None

    name: str

    block_size: int = 512

    capacity: int

    serial_num = None

    def __init__(self, name, path, disk_type, block_size, capacity, dev_id, index):
        self.name = name
        self.disk_type = disk_type
        self.path = path
        if block_size is not None:
            self.block_size = block_size
        self.capacity = capacity
        self.serial_num = dev_id
        self.index = index

    def erase(self, method, progressbar, progress_label) -> threading.Thread:
        t = threading.Thread(target=self.__erasing, args=(method, progressbar,progress_label))
        t.start()
        return t

    def __erasing(self, method, progressbar, progress_label):
        write_size = self.block_size*32
        with open(self.path, "wb") as drive:
            for i in range(erasing_methods[method]["num_passes"]):
                drive.seek(0)
                prev_percent = 0
                len_write = 0
                progressbar.set(prev_percent)
                print(i)
                while len_write < self.capacity:
                    if erasing_methods[method]["data"][i] == "random":
                        data = get_random_bytes(write_size)
                    else:
                        data = erasing_methods[method]["data"][i] * write_size
                    drive.write(data)
                    len_write += write_size
                    cur_percent = len_write/self.capacity
                    if cur_percent - prev_percent > 0.01:
                        progressbar.set(cur_percent)
                        progress_label.configure(text=f"{int(cur_percent*100)} %")
                        prev_percent = cur_percent