import asyncio
import logging
import time

import websockets
from pojagi_dsp.channel.ecg.generator.wavetable import ECGWaveTableSynthesizer
from pojagi_dsp.channel.ecg.generator.wavetable.sinus import (
    SinusWaveTable, TachycardiaWaveTable)

if __name__ != "__main__":
    raise ImportWarning("This script is not intended to be imported.")

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
log = logging.getLogger(__name__)

PORT = 7890


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


async def producer_handler(websocket: websockets.WebSocketServerProtocol):
    srate = 50

    ecg = ECGWaveTableSynthesizer(
        tables={
            (0, 90): SinusWaveTable(),
            (70, 300): TachycardiaWaveTable(),
        },
        heart_rate=60,
        srate=srate,
    )
    while True:
        try:
            message = next(ecg)
            await websocket.send(str(message))
            await asyncio.sleep(1/srate)
        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)

try:
    asyncio.get_event_loop().run_forever()
except KeyboardInterrupt:
    log.info("exiting...")