import asyncio
from struct import Struct
import click


PORT = 8887
ENC = 'utf-8'

DATA = {}


def perform_GET(key):
    return DATA.get(key, b'')


def perform_SET(key, value):
    DATA[key] = value
    return b'1'


pkt_len = Struct('!I')


async def handle_connection(reader, writer):
    try:
        while True:
            cmdlen = pkt_len.unpack(await reader.readexactly(pkt_len.size))[0]
            data = await reader.readexactly(cmdlen)
            cmd, params = data.split(b' ', 1)

            if cmd == b'GET':
                result = perform_GET(params)
            elif cmd == b'SET':
                result = perform_SET(*params.split(b' ', 1))
            else:
                result = b'ERROR'

            writer.write(pkt_len.pack(len(result)))
            writer.write(result)
            await writer.drain()
    except asyncio.streams.IncompleteReadError:
        writer.close()


async def issue_client(loop, cmd):
    reader, writer = await asyncio.open_connection('127.0.0.1', PORT, loop=loop)

    writer.write(pkt_len.pack(len(cmd)))
    writer.write(cmd)

    datalen = pkt_len.unpack(await reader.readexactly(pkt_len.size))[0]
    data = await reader.readexactly(datalen)

    writer.close()
    return data


@click.group()
def cli():
    pass


@cli.command()
@click.argument('cmd', nargs=-1)
def client(cmd):
    loop = asyncio.get_event_loop()
    result = loop.run_until_complete(issue_client(loop, ' '.join(cmd).encode(ENC)))
    loop.close()
    click.echo(result.decode(ENC))


@cli.command()
def server():
    loop = asyncio.get_event_loop()
    coro = asyncio.start_server(handle_connection, '0.0.0.0', PORT, loop=loop)
    server = loop.run_until_complete(coro)

    # Serve requests until Ctrl+C is pressed
    print('Serving on {}'.format(server.sockets[0].getsockname()))
    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass

    # Close the server
    server.close()
    loop.run_until_complete(server.wait_closed())
    loop.close()


if __name__ == '__main__':
    cli()