from asyncio import AbstractEventLoop
from infra.log.module_logger import ModuleLogger
import json
import asyncio
from .async_client import AsyncMQClient


class AsyncMQSubscriber(AsyncMQClient):
    def __init__(self, channel_name: str) -> None:
        super().__init__(channel_name=channel_name)
        self.process_callable = None
        self.routing_key = self.channel_name
        self.consumer_callbacks = {}
        self.consumer_callbacks_lock = asyncio.Lock()  # Async lock for async context
        self.module_logger = ModuleLogger(sender_id="AsyncMQSubscriber")

    async def process_incoming_message(self, message):
        """Processing incoming message from RabbitMQ"""
        await message.ack()
        body = message.body
        if body:
            async with self.consumer_callbacks_lock:  # Use async lock for safe concurrent access
                for registry_key, callback_info in self.consumer_callbacks.items():
                    try:
                        await callback_info["method"](
                            registry_key, json.loads(body), callback_info["args"]
                        )
                    except Exception as err:
                        # Log each exception that occurs within callback processing
                        await self.module_logger.log_exception(
                            exception=err,
                            text=f"Error processing message for consumer '{registry_key}'",
                        )

    async def subscribe(self, max_retries=10, event_loop: AbstractEventLoop = None):
        """Attempts to bind and consume messages, with retry mechanism."""
        retries = 0
        while retries < max_retries:
            try:
                await self.bind(max_retries=5, event_loop=event_loop)
                await self.queue.consume(
                    no_ack=False, exclusive=True, callback=self.process_incoming_message
                )
                break  # Exit loop if subscription is successful
            except Exception as e:
                await self.module_logger.log_exception(
                    exception=e,
                    text=f"Failed to subscribe at {retries} time, will retry",
                )
                retries += 1
                await asyncio.sleep(5)  # Delay before retrying
        else:
            await self.module_logger.log_exception(
                exception=ConnectionError(
                    f"Exceeded max retries ({max_retries}) for subscription."
                ),
                text=f"Subscription failed for {self.channel_name} after {max_retries} attempts.",
            )

    async def register_consumer(
        self,
        registry_key: str,  # a unique string to identify the callback
        callback_method,
        args: dict,
    ):
        """Register a consumer callback with a unique key."""
        async with self.consumer_callbacks_lock:
            self.consumer_callbacks[registry_key] = {
                "method": callback_method,
                "args": args,
            }

    async def unregister_consumer(
        self,
        registry_key: str,  # a unique string to identify the callback
    ):
        """Unregister a consumer callback by its key."""
        async with self.consumer_callbacks_lock:
            if registry_key in self.consumer_callbacks:
                del self.consumer_callbacks[registry_key]

    async def clear_all_consumers(self):
        """Unregister all consumer callbacks."""
        async with self.consumer_callbacks_lock:
            self.consumer_callbacks.clear()