import asyncio
import contextvars
import sys

cvar1 = contextvars.ContextVar("cvar1")
cvar2 = contextvars.ContextVar("cvar2")
cvar3 = contextvars.ContextVar("cvar3")


def print_diagnostics(label):
    task = t.get_name() if (t := asyncio.current_task()) else None
    context = {c.name: v for c, v in contextvars.copy_context().items()}
    print(f"{label}: {task=}, {context=}\n{'-' * 80}")


class DemoProtocol(asyncio.Protocol):
    def __init__(self, on_conn_lost):
        self.transport = None
        self.on_conn_lost = on_conn_lost
        self.tasks = set()

    def connection_made(self, transport):
        print_diagnostics("connection_made")

        self.transport = transport

    def data_received(self, data):
        print_diagnostics("data_received")

        task = asyncio.create_task(self.asgi())
        self.tasks.add(task)
        task.add_done_callback(self.tasks.discard)

        self.transport.pause_reading()

    def connection_lost(self, exc):
        print_diagnostics("connection_lost")
        if not self.on_conn_lost.done():
            self.on_conn_lost.set_result(True)

    async def asgi(self):
        print_diagnostics("asgi start")

        cvar1.set(True)

        # make sure that we only resume after the pause
        # otherwise the resume does nothing
        while not self.transport._paused:
            await asyncio.sleep(0.1)

        cvar2.set(True)

        self.transport.resume_reading()

        cvar3.set(True)

        print_diagnostics("asgi end")


async def main():
    print(f"Python: {sys.version}\n{'-' * 80}")

    loop = asyncio.get_running_loop()
    on_conn_lost = loop.create_future()

    host, port = "127.0.0.1", 8888

    async with await loop.create_server(lambda: DemoProtocol(on_conn_lost), host, port):
        reader, writer = await asyncio.open_connection(host, port)
        writer.write(b"anything")
        await writer.drain()
        writer.close()
        await writer.wait_closed()
        await on_conn_lost


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