import structio
import sys
import time

_print = print


def print(*args, **kwargs):
    sys.stdout.write(f"[{time.strftime('%H:%M:%S')}] ")
    _print(*args, **kwargs)


async def test(host: str, port: int, bufsize: int = 4096, keepalive: bool = False):
    print(f"Attempting a connection to {host}:{port} {'in keep-alive mode' if keepalive else ''}")
    socket = await structio.socket.connect_tcp_ssl_socket(host, port)
    buffer = b""
    print("Connected")
    # Ensures the code below doesn't run for more than 5 seconds
    with structio.skip_after(5) as scope:
        # Closes the socket automatically
        async with socket:
            print("Entered socket context manager, sending HTTP request data")
            await socket.send_all(
                f"GET / HTTP/1.1\r\nUser-Agent: Structio/0.1.0\r\nAccept: */*\r\nHost: {host}"
                f"\r\nAccept-Encoding: gzip, deflate, br\r\nConnection: {'close' if not keepalive else 'keep-alive'}"
                f"\r\n\r\n".encode()
            )
            print("Data sent")
            while True:
                # We purposely do NOT check for the end of the response (\r\n) so that when the
                # connection is in keep-alive mode we hang and let our timeout expire the whole
                # block
                print(
                    f"Requesting up to {bufsize} bytes (current response size: {len(buffer)})"
                )
                data = await socket.receive(bufsize)
                if data:
                    print(f"Received {len(data)} bytes")
                    buffer += data
                else:
                    print("Received empty stream, closing connection")
                    break
    if buffer:
        data = buffer.decode().split("\r\n")
        print(
            f"HTTP Response below {'(might be incomplete)' if scope.timed_out else ''}:"
        )
        _print(f"Response: {data[0]}")
        _print("Headers:")
        content = False
        for i, element in enumerate(data):
            if i == 0:
                continue
            else:
                if not element.strip() and not content:
                    sys.stdout.write("\nContent:")
                    content = True
                if not content:
                    _print(f"\t{element}")
                else:
                    for line in element.split("\n"):
                        _print(f"\t{line}")
    _print("Done!")


structio.run(test, "google.com", 443, 256)
# With keep-alive on, our timeout will kick in
structio.run(test, "google.com", 443, 256, True)