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


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


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

    if cmd == 'get':
        make_get_request(args)

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

        src, dest = paths

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

        if status == 0:
            with open(dest, 'wb') as f:
                f.write(db[src])

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


def make_get_request(path):
    packet = pack_data(json.dumps({'method': 'get', 'path': path}))
    sock.sendall(packet)
    response = get_response()
    return handle_response(response)


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


def get_response():
    buffer = b''
    while len(buffer) < DATATYPE_SIZE:
        data = sock.recv(DATATYPE_SIZE - len(buffer))
        if not data:
            raise ConnectionError('Connection was closed')

        buffer += data

    size_packed, buffer = buffer[:DATATYPE_SIZE], buffer[DATATYPE_SIZE:]
    if not size_packed:
        raise ConnectionError('Connection was closed')

    (size,) = struct.unpack('I', size_packed)
    data = sock.recv(size - len(buffer))
    return parse_response(buffer + data)


def download_file(filesize):
    data = b''

    if TQDM:
        for count in tqdm(range(0, filesize, CHUNK_SIZE)):
            (size,) = struct.unpack('I', sock.recv(DATATYPE_SIZE))
            data += sock.recv(size)
    else:
        count = 0
        while count < filesize:
            (size,) = struct.unpack('I', sock.recv(DATATYPE_SIZE))
            data += sock.recv(size)
            count += size

    return data


def handle_response(data):
    if data['status'] == 0:
        print(f'Загружаем файл размером {data["size"]} байт...')
        file_data = download_file(data['size'])
        db[data['path']] = file_data
        print('Загрузка завершена')

    elif data['status'] == 1:
        print(data['reason'])

    elif data['status'] == 2:
        print('Неизвестная ошибка:')
        print(data['reason'])

    return data['status']


def try_connect():
    cnt = 0
    sock = socket.socket()
    while cnt < 10:
        try:
            cnt += 1
            sock.connect(ADDR)

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

        else:
            return sock

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


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

sock = try_connect()
if sock is None:
    quit()


db = shelve.open('backdoor.db')

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

    except KeyboardInterrupt:
        break

sock.close()