import time
import websockets
import asyncio
from pojagi_dsp.channel.ecg.generator.wavetable import ECGWaveTableSynthesizer
from pojagi_dsp.channel.ecg.generator.wavetable.sinus import SinusWaveTable, TachycardiaWaveTable

PORT = 7890

async def consumer_handler(websocket):
    async for message in websocket:
        print(f"message received: {message}")

async def producer_handler(websocket):
    ecg = ECGWaveTableSynthesizer(
        tables={
            (0,90): SinusWaveTable(),
            (70, 300): TachycardiaWaveTable(),
        },
        heart_rate=70,
        srate=50,
    )
    while True:
        try:
            message = next(ecg)
            await websocket.send(str(message))
            time.sleep(0.01)
        except websockets.exceptions.ConnectionClosed as e:
            print("A client just disconnected")
            break

async def handler(websocket, path):
    while True:
        print(f"New connection. Path: {path}")
        consumer_task = asyncio.create_task(consumer_handler(websocket))
        producer_task = asyncio.create_task(producer_handler(websocket))
        done, pending = await asyncio.wait(
            [consumer_task, producer_task],
            return_when=asyncio.FIRST_COMPLETED,
        )
        for task in pending:
            task.cancel()

# Start the server
start_server = websockets.serve(handler, "0.0.0.0", PORT)
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()