import os
from typing import Optional

import pika
from testcontainers.core.container import DockerContainer
from testcontainers.core.waiting_utils import wait_container_is_ready


class RabbitMqContainer(DockerContainer):
    """
    Test container for RabbitMQ.

    Example
    -------
    The example spins up a RabbitMQ broker and uses the `pika` client library
    (https://pypi.org/project/pika/) establish a connection to the broker.
    ::
        from testcontainer.rabbitmq import RabbitMqContainer
        import pika

        with RabbitMqContainer("rabbitmq:3.9.10") as rabbitmq:

            connection = pika.BlockingConnection(rabbitmq.get_connection_params())
            channel = connection.channel()
    """

    RABBITMQ_NODE_PORT = os.environ.get("RABBITMQ_NODE_PORT", 5672)
    RABBITMQ_DEFAULT_USER = os.environ.get("RABBITMQ_DEFAULT_USER", "guest")
    RABBITMQ_DEFAULT_PASS = os.environ.get("RABBITMQ_DEFAULT_PASS", "guest")

    def __init__(
        self,
        image: str = "rabbitmq:latest",
        port: Optional[int] = None,
        username: Optional[str] = None,
        password: Optional[str] = None,
    ) -> None:
        """Initialize the RabbitMQ test container.

        Args:
            image (str, optional):
                The docker image from docker hub. Defaults to "rabbitmq:latest".
            port (int, optional):
                The port to reach the AMQP API. Defaults to 5672.
            username (str, optional):
                Overwrite the default username which is "guest".
            password (str, optional):
                Overwrite the default username which is "guest".
        """
        super(RabbitMqContainer, self).__init__(image=image)
        self.RABBITMQ_NODE_PORT = port or int(self.RABBITMQ_NODE_PORT)
        self.RABBITMQ_DEFAULT_USER = username or self.RABBITMQ_DEFAULT_USER
        self.RABBITMQ_DEFAULT_PASS = password or self.RABBITMQ_DEFAULT_PASS

        self.with_exposed_ports(self.RABBITMQ_NODE_PORT)
        self.with_env("RABBITMQ_NODE_PORT", self.RABBITMQ_NODE_PORT)
        self.with_env("RABBITMQ_DEFAULT_USER", self.RABBITMQ_DEFAULT_USER)
        self.with_env("RABBITMQ_DEFAULT_PASS", self.RABBITMQ_DEFAULT_PASS)

    @wait_container_is_ready(pika.exceptions.IncompatibleProtocolError)
    def readiness_probe(self) -> bool:
        """Test if the RabbitMQ broker is ready."""
        connection = pika.BlockingConnection(self.get_connection_params())
        if connection.is_open:
            connection.close()
            return self
        raise RuntimeError("Could not open connection to RabbitMQ broker.")

    def get_connection_params(self) -> pika.ConnectionParameters:
        """
        Get connection params as a pika.ConnectionParameters object.
        For more details see:
        https://pika.readthedocs.io/en/latest/modules/parameters.html
        """
        credentials = pika.PlainCredentials(username=self.RABBITMQ_DEFAULT_USER,
                                            password=self.RABBITMQ_DEFAULT_PASS)

        return pika.ConnectionParameters(
            host=self.get_container_host_ip(),
            port=self.get_exposed_port(self.RABBITMQ_NODE_PORT),
            credentials=credentials,
        )

    def start(self):
        """Start the test container."""
        super().start()
        self.readiness_probe()
        return self
