"""
Websockets client for micropython
Based very heavily off
https://github.com/aaugustin/websockets/blob/master/websockets/client.py
"""

import binascii as binascii
import random as random
# import ussl
import uasyncio

from uws_protocol import Websocket, urlparse

# TODO: Mick: Commented out because not needed?
# context = ussl.create_default_context()


class WebsocketClient(Websocket):
    is_client = True


async def connect(uri):
    """
    Connect a websocket.
    """

    uri = urlparse(uri)
    assert uri

    print("open connection %s:%s" % (uri.hostname, uri.port))

    # reader, writer = yield from uasyncio.wait_for(uasyncio.open_connection(uri.hostname, uri.port), timeout=5)
    reader, writer = await uasyncio.open_connection(uri.hostname, uri.port)

    print("opened asyncio streams: (%s, %s)" % (reader, writer))

    if uri.protocol == "wss":
        # sock = ussl.wrap_socket(sock)
        raise Exception("SSL Sockets not supported using uasyncio")

    # Sec-WebSocket-Key is 16 bytes of random base64 encoded
    key = binascii.b2a_base64(bytes(random.getrandbits(8) for _ in range(16)))[:-1]

    headers = [
        "GET %s HTTP/1.1" % (uri.path or "/"),
        "Host: %s:%s" % (uri.hostname, uri.port),
        "Connection: Upgrade",
        "Upgrade: websocket",
        "Sec-WebSocket-Key: %s" % key.decode("UTF-8"),
        "Sec-WebSocket-Version: 13",
        "Origin: http://{hostname}:{port}".format(hostname=uri.hostname, port=uri.port)
    ]
    for header in headers:
        # print(header)
        await writer.awrite(header + "\r\n")
        await writer.drain()

    await writer.awrite("\r\n")
    await writer.drain()

    print("Sent headers, waiting for response")
    header = await reader.readline()
    print("Headers received:\r\n%s" % header.decode("UTF-8"))
    assert header.startswith(b"HTTP/1.1 101 "), header
    header = header[:-2]

    # We don't (currently) need these headers
    # FIXME: should we check the return key?
    while header:
        header = await reader.readline()
        print("%s" % header.decode("UTF-8")[:-2])
        if not header or header.decode("UTF-8")[:2] == "\r\n":
            print("End of HTTP response")
            break

    return WebsocketClient(reader, writer)