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


def parse_request(client_socket: socket.socket):
    request_header_size = struct.unpack("!Q", read_bytes(client_socket, 8))[0]
    request_map = json.JSONDecoder().decode(read_bytes(client_socket, request_header_size).decode("utf-8"))
    return request_map


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(request_info, socket_to_server):
    full_request = generate_request(request_info)

    socket_to_server.sendall(full_request)

    if request_info['action'] == 'end' or request_info['action'] == 'start':
        return

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

    return responseJson


def receive_response(server_socket):
    request_map = parse_request(server_socket)
    print("receiving response:\n" + json.dumps(request_map, ensure_ascii=False))
    with open('result.json', 'w', encoding='utf-8') as f:
        json.dump(request_map, f, ensure_ascii=False, indent=4)


if __name__ == '__main__':
    # 使用方法 python .\connect.py --ip 127.0.0.1 --port 7777
    # crawling --word computer --cookie 95f94e1ab71bdf96b85fef6e8f746c58eeb5f9fa --pages_start 1 --pages_end 10
    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_crawling.add_argument('--pages_end', type=int, required=True)
    parser_crawling.add_argument('--pages_start', type=int, required=True)
    parser_crawling.add_argument('--cookie', 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()

    local_port = 9010
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.connect((args.ip, int(args.port)))
    request = {'action': 'start'}
    send_request(request, server_socket)

    if args.action == 'crawling':
        request = {'action': 'crawl zhiwang', 'word': args.word, 'pages_start': args.pages_start,
                   'pages_end': args.pages_end, 'cookie': args.cookie}
    elif args.action == 'login' or args.action == 'register':
        request = {'action': args.action, 'user': args.user, 'password': args.password}
    response = send_request(request, server_socket)
    print(response['cookie'])

    if args.action == 'crawling':
        receive_response(server_socket)

    request = {'action': 'end'}
    send_request(request, server_socket)
    server_socket.close()