from dataclasses import dataclass

import trio
from itertools import count
from functools import partial

CONNECTION_COUNTER = count()
TASK_COUNTER = count()


@dataclass
class Operation:
    id: int
    sleep_s: int


class Database:
    ops: list[Operation]
    id_counter: int = 1

    def __init__(self):
        self.ops = []


@dataclass
class DatabaseHandler:
    db: Database
    lock: trio.Lock = trio.Lock()


async def connection_handler(dbh: DatabaseHandler, stream):
    ident = next(CONNECTION_COUNTER)
    async for data in stream:
        print(f"{ident}: received data {data!r}")
        await stream.send_all(data)


async def op_execution(op: Operation):
    task_id = next(TASK_COUNTER)
    print(f"Joe {op.id} {task_id} - Started sleeping for {op.sleep_s}sec...")
    await trio.sleep(op.sleep_s)
    print(f"Joe {op.id} {task_id} - Slept well!")


async def op_spawner(dbh: DatabaseHandler, nursery: trio.Nursery):
    while True:
        async with dbh.lock:
            for op in dbh.db.ops:
                nursery.start_soon(op_execution, op)
        await trio.sleep(10)


async def main():
    db = Database()
    for _ in range(5):
        i = db.id_counter
        db.id_counter += 1
        db.ops.append(Operation(i, i + 1))
    dbh = DatabaseHandler(db)
    async with trio.open_nursery() as nursery:
        nursery.start_soon(op_spawner, dbh, nursery)
        await trio.serve_tcp(partial(connection_handler, dbh), 3000, host='0.0.0.0')


if __name__ == '__main__':
    trio.run(main)