import asyncio
import os
from enum import Enum

from typing import Any

from pydantic import BaseModel, Field
from faststream import FastStream
from faststream.nats import NatsBroker

from models import ProbeRegister, ProbeUnregister, ProbeHeartBeat

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

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



@broker.subscriber(f"discovery.probe")
async def probe_discovery(data: ProbeRegister | ProbeHeartBeat | ProbeUnregister):
    print(20*'-')
    print(type(data), "->", data)
    return True

async def main():
    print("START")
    await app.run()

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