# -*- coding: utf-8 -*-
"""
Created on Tue Dec 10 14:54:00 2021

@author: yqlim
"""
from ast import Call
from typing import Callable, Dict, List, Union
import pika
import json
import time

import logging
logger = logging.getLogger(__name__)

class RabbitMQ(object):
    def __init__(self) -> None:
        pass

    @classmethod
    def make_task_basic(cls, task_name: str, data: Dict) -> str:
        data['task_name'] = task_name
        return json.dumps(data)
    
    @classmethod
    def join_queue_basic(cls, queue_name: str, body: Dict, host: str = 'localhost',
                            retries: int = 10) -> None:
        for i in range(retries):
            try:
                connection = pika.BlockingConnection(pika.ConnectionParameters(host))
                channel = connection.channel()
                channel.queue_declare(queue=queue_name)
                channel.basic_publish(exchange='', routing_key=queue_name, body=body)
                connection.close()
                break
            except Exception as e:
                msg = 'error at queue ({queue_name}) - {error_message}\nconnection closed...and restarted'.format(
                            queue_name=queue_name, error_message=e)
                logger.error(msg)   
                continue             

    @classmethod
    def consume_basic(cls, queue_name: str, callback: Callable, auto_acknowledge: bool = True,
                            host: str='localhost', heartbeat: int=900):
        while True:
            try:
                connection = pika.BlockingConnection(pika.ConnectionParameters(host=host, heartbeat=heartbeat))
                channel = connection.channel()
                channel.queue_declare(queue=queue_name)
                channel.basic_qos(prefetch_count=1)
                channel.basic_consume(queue=queue_name, auto_ack=auto_acknowledge, 
                                        on_message_callback=callback)
                channel.start_consuming()
            except Exception as e:
                msg = 'error at queue ({queue_name}) - {error_message}\nconnection closed...and restarted'.format(
                            queue_name=queue_name, error_message=e)
                logger.error(msg)
                continue

    @classmethod
    def join_fanout_exchange(cls, exchange_name: str, body: Dict, 
        passive: bool = False, durable: bool = False, auto_delete: bool = True,
        internal: bool = False, arguments: Union[Dict, None] = None, 
        callable: Union[pika.frame.Method, None] = None, host: str = 'localhost'):
        
        connection = pika.BlockingConnection(pika.ConnectionParameters(host))
        channel = connection.channel()
        channel.exchange_declare(
            exchange = exchange_name,
            exchange_type = 'fanout',
            passive = passive,
            durable = durable,
            auto_delete = auto_delete,
            internal = internal,
            arguments = arguments,
            callable = callable
        )
        channel.basic_publish(exchange=exchange_name, routing_key='', body=body)
        connection.close()

    @classmethod
    def consume_fanout(cls, exchange_name: str, callback: Callable, auto_acknowledge: bool = True,
                            host: str = 'localhost', heartbeat: int = 900):
        while True:
            try:
                connection = pika.BlockingConnection(pika.ConnectionParameters(host=host, heartbeat=heartbeat))
                channel = connection.channel()
                channel.exchange_declare(exchange=exchange_name, exchange_type='fanout')
                result = channel.queue_declare(queue='', exclusive=True)
                queue_name = result.method.queue

                channel.queue_bind(exchange=exchange_name, queue=queue_name)
                
                channel.basic_qos(prefetch_count=1)
                channel.basic_consume(queue=queue_name, auto_ack=auto_acknowledge, 
                                        on_message_callback=callback)
                channel.start_consuming()
            except Exception as e:
                msg = 'error at queue ({queue_name}) - {error_message}\nconnection closed...and restarted'.format(
                            queue_name=queue_name, error_message=e)
                logger.error(msg)
                continue        

    @classmethod
    def join_direct_exchange(cls, exchange_name: str, routing_key: str, body: Dict, 
        passive: bool = False, durable: bool = False, auto_delete: bool = False,
        internal: bool = False, arguments: Union[Dict, None] = None, 
        callable: Union[pika.frame.Method, None] = None, host: str = 'localhost'):
        
        connection = pika.BlockingConnection(pika.ConnectionParameters(host))
        channel = connection.channel()
        channel.exchange_declare(
            exchange = exchange_name,
            exchange_type = 'direct',
            passive = passive,
            durable = durable,
            auto_delete = auto_delete,
            internal = internal,
            arguments = arguments
        )
        channel.basic_publish(exchange=exchange_name, routing_key=routing_key, body=body)
        connection.close()

    @classmethod
    def consume_direct(cls, exchange_name: str, routing_keys: List[str], callback: Callable, auto_acknowledge: bool = True,
                            host: str = 'localhost', heartbeat: int = 900):
        while True:
            try:
                connection = pika.BlockingConnection(pika.ConnectionParameters(host=host, heartbeat=heartbeat))
                channel = connection.channel()
                channel.exchange_declare(exchange=exchange_name, exchange_type='direct')
                result = channel.queue_declare(queue='', exclusive=True)
                queue_name = result.method.queue
                
                for routing_key in routing_keys:
                    channel.queue_bind(exchange=exchange_name, queue=queue_name, routing_key=routing_key)
                
                channel.basic_qos(prefetch_count=1)
                channel.basic_consume(queue=queue_name, auto_ack=auto_acknowledge, 
                                        on_message_callback=callback)
                channel.start_consuming()
            except Exception as e:
                msg = 'error at queue ({queue_name}) - {error_message}\nconnection closed...and restarted'.format(
                            queue_name=queue_name, error_message=e)
                logger.error(msg)
                continue   

class RabbitCallback(object):
    def __init__(self, func: Callable) -> None:
        self.func = func 

        self.ch = None
        self.method = None
        self.properties = None
        self.body = None

    def consume(self, ch: pika.adapters.blocking_connection.BlockingChannel, 
                     method: pika.spec.Basic.Deliver, 
                     properties: pika.spec.BasicProperties, 
                     body: bytes) -> None:  
        self.ch = ch
        self.method = method
        self.properties = properties
        self.body = json.loads(body)

        self.func(self.body)
