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

def generate_request(request) -> 'bytes':
    """
    根据传入的dict生成请求
    请求包含 8字节头长度+头数据
    :param request: dict
    :return: bytes 请求数据
    """
    request_bytes = JSONEncoder().encode(request).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 crawling(ip, port, word: str):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP) as socket_to_server:
        socket_to_server.connect((ip, int(port)))
        request = dict()
        request['action'] = 'crawl zhiwang'
        request['word'] = word

        full_request = generate_request(request)

        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['crawl zhiwang']


if __name__ == '__main__':
    parser = argparse.ArgumentParser('Automanager')
    parser.add_argument('--ip', type = str, required = True)
    parser.add_argument('--port', type = str, required = True)
    parser.add_argument('--word', type = str, required = True)

    args = parser.parse_args()

    crawling(args.ip, args.port, args.word)