import pika
import os
import json
import dotenv
import threading
import time

dotenv.load_dotenv()

RABBITMQ_URL = os.getenv("RABBITMQ_URL")

def send_topic_extraction_request(payload: dict):
    """Simulate ml_server sending a topic extraction request to gpu_server."""
    params = pika.URLParameters(RABBITMQ_URL)
    params.heartbeat = 5
    params.blocked_connection_timeout = 2
    connection = pika.BlockingConnection(params)
    channel = connection.channel()
    channel.queue_declare(queue="gpu_server", durable=True)
    
    message = json.dumps(payload).encode("utf-8")
    channel.basic_publish(
        exchange="",
        routing_key="gpu_server",
        body=message,
        properties=pika.BasicProperties(delivery_mode=2)
    )
    print("Topic extraction request sent from ml_server to gpu_server.")
    connection.close()

def listen_for_results():
    """Simulate ml_server listening for topic extraction results on ml_server queue."""
    params = pika.URLParameters(RABBITMQ_URL)
    params.heartbeat = 5
    params.blocked_connection_timeout = 2
    connection = pika.BlockingConnection(params)
    channel = connection.channel()
    channel.queue_declare(queue="ml_server", durable=True)
    
    def callback(ch, method, properties, body):
        try:
            result = json.loads(body)
            print("Received topic extraction result:")
            print(json.dumps(result, indent=2))
            ch.basic_ack(delivery_tag=method.delivery_tag)
        except Exception as e:
            print("Error processing message:", e)
            ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
    
    channel.basic_consume(queue="ml_server", on_message_callback=callback)
    print("Listening for topic extraction results on ml_server queue...")
    channel.start_consuming()

if __name__ == "__main__":
    payload = {
        "pattern": "topic_extraction",
        "data": {
            "input_files": [
            {
                "key": "",
                "url": "https://qualifications.pearson.com/content/dam/pdf/A%20Level/Mathematics/2017/specification-and-sample-assesment/a-level-l3-mathematics-specification-issue4.pdf",
                "type": "specification",
                "page": [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 34, 35, 36, 37, 38, 39, 40, 41, 42]
            }
            ],
            "topics": [
            {
                "title": "",
                "id": 123
            }
            ]
        }
    }
    
    producer_thread = threading.Thread(target=send_topic_extraction_request, args=(payload,))
    producer_thread.start()
    
    time.sleep(1)
    
    listen_for_results()