import socket
import struct
import argparse
from json import JSONEncoder, JSONDecoder


def generate_request(request_info) -> 'bytes':
    """
    根据传入的dict生成请求
    请求包含 8字节头长度+头数据
    :param request_info: dict
    :return: bytes 请求数据
    """
    request_bytes = JSONEncoder().encode(request_info).encode("utf-8")
    return struct.pack("!Q", len(request_bytes)) + request_bytes


def read_bytes(s: 'socket.socket', size: 'int') -> 'bytes':
    """
    从socket读取size个字节
    :param s:套接字
    :param size:要读取的大小
    :return:读取的字节数，在遇到套接字关闭的情况下，返回的数据的长度可能小于 size
    """
    data = ''.encode('utf-8')
    while len(data) < size:
        rsp_data = s.recv(size - len(data))
        data += rsp_data
        if len(rsp_data) == 0:
            break
    return data


def send_request(ip, port, request_info):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP) as socket_to_server:
        socket_to_server.connect((ip, int(port)))

        full_request = generate_request(request_info)

        socket_to_server.sendall(full_request)

        responseJson = JSONDecoder().decode(
            read_bytes(socket_to_server, struct.unpack('!Q', socket_to_server.recv(8))[0]).decode(
                "utf-8"))

        return responseJson[request_info['action']]


if __name__ == '__main__':
    # 使用方法 python.exe .\connect.py --ip 127.0.0.1 --port 7777 login --user 1 --password 1
    parser = argparse.ArgumentParser('connect-manager')
    parser.add_argument('--ip', type=str, required=True)
    parser.add_argument('--port', type=str, required=True)
    subparsers = parser.add_subparsers(help='provide actions including crawling, login, register',
                                       dest='action')  # 创建子解析器

    parser_crawling = subparsers.add_parser('crawling')
    parser_crawling.add_argument('--word', type=str, required=True)

    parser_login = subparsers.add_parser('login')
    parser_login.add_argument('--user', type=str, required=True)
    parser_login.add_argument('--password', type=str, required=True)

    parser_register = subparsers.add_parser('register')
    parser_register.add_argument('--user', type=str, required=True)
    parser_register.add_argument('--password', type=str, required=True)

    args = parser.parse_args()

    request = dict()

    if args.action == 'crawling':
        request = {'action': 'crawl zhiwang', 'word': args.word}
    elif args.action == 'login' or args.action == 'register':
        request = {'action': args.action, 'user': args.user, 'password': args.password}
    response = send_request(args.ip, args.port, request)
    print(response)