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

SUBJECT="mtr_client.discovery"

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)

class JobCreated(BaseModel):
    job_id: str = Field(..., description="Job identifier")
    name: str = Field(..., description="Job name")


class JobResponse(BaseModel):
    job_id: str = Field(..., description="Job identifier")
    status: str = Field(..., description="Job status")

class ProbeDiscover(BaseModel):
    name: str = Field(..., description="Probe identifier")
    type: str = Field(..., description="Probe name")

class ProbeCommandTypeEnum(Enum):
    ADD_PROBE = "add probe"
    REMOVE_PROBE = "remove probe"

class ProbeTypeEnum(Enum):
    OTHER = 'other'
    TAS_SERVER = 'tas server'

class DiscoverApiResponseModel(BaseModel):
    probe_type: ProbeTypeEnum
    command: ProbeCommandTypeEnum
    name: str


@broker.subscriber("test_topic.publish")
async def handle_job(event: Any):
    print(f"{SUBJECT}.broadcast received job_id={event}")

    #response_msg = JobCreated(job_id=event.job_id, name="created1").model_dump()
    print(f"{event=}")
    print(f"{ProbeTypeEnum.TAS_SERVER=}")
    #response_msg = DiscoverApiResponseModel(name="TAS_ZAG", probe_type=ProbeTypeEnum.TAS_SERVER.value, command=ProbeCommandTypeEnum.REMOVE_NODE.value).model_dump()
    response_msg = DiscoverApiResponseModel(name="TAS_BUD", probe_type=ProbeTypeEnum.TAS_SERVER, command=ProbeCommandTypeEnum.ADD_PROBE).model_dump()
    print(response_msg)
    #print(f"{response_msg=}")
    #response_msg = {'name': 'TAS_ZAG', 'probe_type': 'tas server', 'command': 'add node'}
    print("===", await broker.publish(response_msg, subject=f"test_topic.request"))
    #print("===", await broker.publish("neco", subject="mtr"))
    return f"Hello"

async def main():
    await app.run()

if __name__ == "__main__":
    #SUBJECT=sys.argv[1]
    asyncio.run(main())