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(ip, port, request_info):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP) as socket_to_server:
        socket_to_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        socket_to_server.bind(('', 9005))
        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


def receive_response():
    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.bind(('', 9005))
    server_socket.listen()
    while True:
        client_socket, _ = server_socket.accept()
        request_map = parse_request(client_socket)
        if request_map['type'] == 'response':
            print("receiving response:\n" + json.dumps(request_map, ensure_ascii=False))
            break


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()

    request = dict()

    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(args.ip, args.port, request)
        print(response['cookie'])

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