import asyncio
import os

from faststream.nats import NatsBroker
from faststream.nats.message import NatsMessage

from broadcast.models import ProbeRegister, ProbeHeartBeat, ProbeUnregister

NATS_SERVERS = os.getenv('NATS_URL', 'nats://localhost:4222').split(',')
NATS_TOKEN = os.getenv('NATS_TOKEN', '')

broker = NatsBroker(servers=NATS_SERVERS, token=NATS_TOKEN)

async def main():
    msg_register = ProbeRegister(
        command="register",
        probe_id="prvni",
        capabilities="capabilities",
        probe_type="mtr"
    )

    msg_heart_beat = ProbeHeartBeat(
        command="hartbeat",
        probe_id="prvni",
        status="OK"
    )

    msg_unregister = ProbeUnregister(
        command="unregister",
        probe_id="prvni",
    )

    await broker.connect()
    #response = await broker.publish(msg.model_dump(), subject=f"discovery.probe")

    response: NatsMessage = await broker.request(
        msg_register.model_dump(),
        subject="discovery.probe",
        timeout=3,
    )
    print(f"{await response.decode()=}")

    await asyncio.sleep(5)
    await broker.publish(msg_heart_beat.model_dump(), subject=f"discovery.probe")

    await asyncio.sleep(5)
    await broker.publish(msg_unregister.model_dump(), subject=f"discovery.probe")

    await broker.stop()

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