import asyncio
import socket

import capnp

capnp.remove_import_hook()

bellairs_capnp = capnp.load("bellairs.capnp")

"""
Mock file server so I can test this.
"""


class TestServer(bellairs_capnp.File.Server):
    def __init__(self):
        self.data = "Voop"

    async def size(self, **kwargs):
        return len(self.data)

    async def read(self, offset, len, **kwargs):
        return self.data


async def client(sock):
    # See https://capnproto.github.io/pycapnp/capnp.html?highlight=twopartyclient#capnp.TwoPartyClient
    # AP: I am a bit unsure why the only thing available here
    # is a two-party client, but this works as a hack.
    client = capnp.TwoPartyClient(sock)
    cap = client.bootstrap()
    cap = cap.cast_as(bellairs_capnp.File)
    result = await cap.size()
    size = result.size
    print("Got %d" % size)
    assert size == 4


async def main():
    # Create a UNIX socket pair, because the network is
    # more complex.
    client_end, server_end = socket.socketpair(socket.AF_UNIX)
    # Create AsyncIoStreams, which is really a wrapper around
    # the C++ Capnproto connection. See https://github.com/capnproto/pycapnp/blob/59a639fa977e4a2e19c6cc60b44cbc9926418710/capnp/lib/capnp.pyx#L1314
    client_end = await capnp.AsyncIoStream.create_connection(sock=client_end)
    server_end = await capnp.AsyncIoStream.create_connection(sock=server_end)
    # Create a TwoPartyServer. Better options
    # are available (e.g., see https://github.com/capnproto/pycapnp/blob/master/examples/async_ssl_server.py#L60), but
    # I did this for now. I think we will need to use
    # `AsyncIoStream.create_server` and do some amount
    # of indirection in there
    _ = capnp.TwoPartyServer(server_end, bootstrap=TestServer())
    print("Started file server")
    await client(client_end)


if __name__ == "__main__":
    asyncio.run(capnp.run(main()))