import os
from typing import Optional

import aio_pika
from aio_pika.abc import (
    AbstractChannel,
    AbstractRobustConnection,
)
from connections.keycloak import (
    get_valid_access_token,
)
from connections.local_files import (
    load_client_info,
)


class RabbitMQManager:
    """Singleton manager for RabbitMQ connection and channels."""

    _instance: Optional["RabbitMQManager"] = None
    _connection: Optional[AbstractRobustConnection] = None
    _consumer_channel: Optional[AbstractChannel] = None
    _publisher_channel: Optional[AbstractChannel] = None
    _heartbeat_channel: Optional[AbstractChannel] = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    async def connect(self):
        """Establish the main connection if not already connected."""
        client_info = load_client_info()
        if not client_info:
            return None
        access_token = get_valid_access_token()
        if not access_token:
            return None
        if self._connection is None or self._connection.is_closed:
            print("Creating new RabbitMQ connection...")
            self._connection = await aio_pika.connect_robust(
                host=os.environ["RABBITMQ_HOST"],
                port=int(os.environ["RABBITMQ_PORT"]),
                login=str(client_info["system_id"]),
                password=access_token,
                virtualhost="/",
            )
            print("RabbitMQ connection established")
        return self._connection

    async def get_consumer_channel(self):
        """Get channel for consuming messages."""
        if self._connection:
            if self._consumer_channel is None or self._consumer_channel.is_closed:
                self._consumer_channel = await self._connection.channel()
                await self._consumer_channel.set_qos(prefetch_count=1, global_=True)
                print("Consumer channel created")
                return self._consumer_channel
            return self._consumer_channel
        return None

    async def get_publisher_channel(self):
        """Get channel for publishing regular messages."""
        if self._connection:
            if self._publisher_channel is None or self._publisher_channel.is_closed:
                self._publisher_channel = await self._connection.channel()
                print("Publisher channel created")
            return self._publisher_channel
        return None

    async def get_heartbeat_channel(self):
        """Get channel for heartbeat publishing."""
        if self._connection:
            if self._heartbeat_channel is None or self._heartbeat_channel.is_closed:
                self._heartbeat_channel = await self._connection.channel()
                print("Heartbeat channel created")
            return self._heartbeat_channel
        return None

    async def close(self):
        """Close all channels and the main connection gracefully."""
        print("Closing RabbitMQ channels and connection...")

        for channel in [
            self._consumer_channel,
            self._publisher_channel,
            self._heartbeat_channel,
        ]:
            if channel and not channel.is_closed:
                await channel.close()

        if self._connection and not self._connection.is_closed:
            await self._connection.close()
            print("RabbitMQ connection closed")


# Create global singleton instance
rabbitmq_manager = RabbitMQManager()