import asyncio
import json
import toml
from loguru import logger
import nats
from providers import get_provider

# Загрузка конфигурации
config = toml.load("config.toml")

NATS_URL = config["nats"]["url"]
INCOMING_TOPIC = config["nats"]["incoming_topic"]
RESPONSE_TOPIC = config["nats"]["response_topic"]

provider = get_provider(
    config["llm"]["provider"],
    api_key=config["llm"]["api_key"],
    model=config["llm"]["model"],
    system_prompt=config["llm"].get("system_prompt", "")
)


async def message_handler(msg):
    logger.info(f"Received message on {msg.subject}: {msg.data}")

    try:
        data = json.loads(msg.data.decode())
        request_id = data.get("request_id")
        user_id = data.get("user_id")
        text = data.get("text")
        if not text:
            logger.warning("Empty text field in message")
            return

        response = await provider.chat([{"role": "user", "content": text}])

        reply = {
            "request_id": request_id,
            "user_id": user_id,
            "response": response,
            "service": "LLM-Service"
        }

        await nc.publish(RESPONSE_TOPIC, json.dumps(reply).encode())
        logger.success(f"Ответ опубликован в {RESPONSE_TOPIC} для request_id: {request_id}")

    except Exception as e:
        logger.exception("Error handling message")


async def main():
    logger.info("Connecting to NATS...")
    global nc
    nc = await nats.connect(servers=[NATS_URL])
    await nc.subscribe(INCOMING_TOPIC, cb=message_handler)
    await nc.flush()
    logger.info(f"Subscribed to {TOPIC_IN}")

    # Просто держим процесс активным
    try:
        logger.info("LLM service is running. Waiting for messages...")
        await asyncio.Event().wait()
    except KeyboardInterrupt:
        logger.info("Shutting down LLM service...")
    finally:
        await nc.drain()


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