"""
MQTT Subscriber — 监听 chongming Worker 的响应

用法:
  python subscriber.py [worker_type] [operation]
  
示例:
  python subscriber.py                    # 监听所有 chongming 响应
  python subscriber.py example add        # 只监听 example.add 的响应
"""

import asyncio
import json
import sys
import logging
from amqtt.client import MQTTClient
from amqtt.mqtt.constants import QOS_1

logging.basicConfig(level=logging.INFO)

PREFIX = "chongming"

# 默认订阅所有 chongming 响应
DEFAULT_TOPIC = f"{PREFIX}/+/+/response"

async def main():
    worker_type = sys.argv[1] if len(sys.argv) > 1 else None
    operation = sys.argv[2] if len(sys.argv) > 2 else None

    if worker_type and operation:
        topic = f"{PREFIX}/{worker_type}/{operation}/response"
    elif worker_type:
        topic = f"{PREFIX}/{worker_type}/+/response"
    else:
        topic = DEFAULT_TOPIC

    client = MQTTClient()
    await client.connect("mqtt://localhost:1883/")

    await client.subscribe([(topic, QOS_1)])
    print(f"[Subscriber] Subscribed to '{topic}'. Waiting for messages... (Press Ctrl+C to stop)")

    try:
        while True:
            message = await client.deliver_message()
            payload = message.data.decode("utf-8")  # type: ignore
            try:
                data = json.loads(payload)
                print(f"[Subscriber] Topic: {message.topic}")  # type: ignore
                print(f"  Response: {json.dumps(data, indent=2, ensure_ascii=False)}")
            except json.JSONDecodeError:
                print(f"[Subscriber] Topic: {message.topic}, Raw: {payload}")  # type: ignore
    except asyncio.CancelledError:
        print("[Subscriber] Stopping...")
    finally:
        await client.unsubscribe([topic])
        await client.disconnect()
        print("[Subscriber] Disconnected")

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("[Subscriber] Interrupted by user")