import socket
from time import sleep
import struct
import json
import shelve

try:
    from tqdm import tqdm
    TQDM = True
except ImportError:
    print('Для вывода прогресса загрузки файлов установите библиотеку tqdm')
    TQDM = False


DATATYPE_SIZE = 4  # int
CHUNK_SIZE = 4096000
HOST, PORT = ADDR = 'localhost', 6000


class Client:
    def __init__(self, server_host, server_port):
        self.server = socket.socket()
        self.db = shelve.open('backdoor.db')

        address = server_host, server_port
        cnt = 0
        while cnt < 10:
            try:
                cnt += 1
                self.server.connect(address)

            except ConnectionRefusedError:
                print('Не удалось подключиться к серверу. Следующая попытка через 3 секунды...')
                sleep(3)

            else:
                break

        else:
            print('Не удалось подключиться к серверу через 10 попыток.\nОстанавливаемся...')
            self.server.close()

    def start(self):
        while True:
            try:
                data = input('Введите команду:\n')
                self.handle_input(data)

            except ValueError as err:
                print(err)

            except KeyboardInterrupt:
                break

        self.server.close()

    def handle_input(self, data):
        cmd, args = data.split(' ', maxsplit=1)

        if cmd == 'get':
            self.download_file(args)

        elif cmd == 'cp':
            paths = args.split()
            if len(paths) != 2:
                raise ValueError('Неверное количество путей в параметрах.')

            src, dest = paths

            # Если файл не присутствует в БД, то предварительно его загружаем
            if src not in self.db:
                self.download_file(src)

            # При любых ошибках при загрузке файла download_file выкинет исключение и исполнение
            # не дойдет до этого блока, поэтому можно не проверять наличие src в БД
            with open(dest, 'wb') as f:
                f.write(self.db[src])

        else:
            raise ValueError('Команда не распознана')

    def download_file(self, path):
        response = self.make_get_request(path)
        if response['status'] == 0:
            print(f'Загружаем файл размером {response["size"]} байт...')

            file_data = self.receive_data(int(response['size']), use_tqdm=True)
            self.db[path] = file_data

            print('Загрузка завершена')

        elif response['status'] == 1:
            raise ValueError(response['reason'])

        elif response['status'] == 2:
            raise ValueError(f'Неизвестная ошибка: {response["reason"]}')

    def make_get_request(self, path):
        packet = Client.pack_str(json.dumps({'method': 'get', 'path': path}))
        self.server.sendall(packet)
        return Client.parse_response(self.receive_data(DATATYPE_SIZE))

    def receive_data(self, data_size, use_tqdm=False):
        data = b''

        # С копипастой надо бы что-то сделать, но я не знаю что
        if not use_tqdm or not TQDM:
            while len(data) < data_size:
                (packet_size,) = struct.unpack('I', self.server.recv(DATATYPE_SIZE))
                data += self.server.recv(packet_size)
                if not data:
                    raise ConnectionError('Connection was closed')

        elif TQDM:
            for count in tqdm(range(0, data_size, CHUNK_SIZE)):
                (packet_size,) = struct.unpack('I', self.server.recv(DATATYPE_SIZE))
                data += self.server.recv(packet_size)
                if not data:
                    raise ConnectionError('Connection was closed')

        return data

    @staticmethod
    def pack_str(data):
        encoded = data.encode()
        return struct.pack('I', len(encoded)) + encoded

    @staticmethod
    def parse_response(data):
        return json.loads(data.decode())


client = Client(HOST, PORT)
client.start()