from prometheus_client import make_asgi_app, CollectorRegistry
from faststream.asgi import AsgiFastStream
from faststream.rabbit import RabbitBroker
from faststream.rabbit.prometheus import RabbitPrometheusMiddleware


registry = CollectorRegistry()


broker = RabbitBroker(middlewares=[RabbitPrometheusMiddleware(registry=registry)])


app = AsgiFastStream(
    broker,
    asyncapi_path="/docs",
    asgi_routes=[
        ("/metrics", make_asgi_app(registry=registry))
    ]
)

from prometheus_client import CollectorRegistry, Counter


class RabbitPrometheusMiddleware(BasePrometheusMiddleware):
    @staticmethod
    def get_broker_name() -> str:
        return "rabbit"


class BasePrometheusMiddleware:
    def __init__(self, registry: CollectorRegistry):
        self._registry = registry
        self._received_messages_counter = Counter(
            name="received_messages",
            documentation="Received messages",
            labelnames=["broker"],
        )
        self._processed_messages_with_error_counter = Counter(
            name="processed_messages_with_error",
            documentation="Processed messages with error",
            labelnames=["broker"],
        )
        self._processed_messages_without_error_counter = Counter(
            name="processed_messages_without_error",
            documentation="Processed messages without error",
            labelnames=["broker"],
        )
        self._registry.register(self._received_messages_counter)
        self._registry.register(self._processed_messages_with_error_counter)
        self._registry.register(self._processed_messages_without_error_counter)

    @staticmethod
    def get_broker_name() -> str:
        raise NotImplementedError

    def __call__(self, msg: Optional[Any]) -> "BaseMiddleware":
        return PrometheusMiddleware(
            msg=msg,
            broker_name=self.get_broker_name(),
            received_messages_counter=self._received_messages_counter,
            processed_messages_with_error_counter=self._processed_messages_with_error_counter,
            processed_messages_without_error_counter=self._processed_messages_without_error_counter,
        )


class PrometheusMiddleware(BaseMiddleware):
    def __init__(
        self,
        msg: Optional[Any] = None,
        *,
        broker_name: str,
        received_messages_counter: "Counter",
        processed_messages_with_error_counter: "Counter",
        processed_messages_without_error_counter: "Counter",
    ) -> None:
        self._broker_name = broker_name
        self._received_messages_counter = received_messages_counter
        self._processed_messages_with_error_counter = processed_messages_with_error_counter
        self._processed_messages_without_error_counter = processed_messages_without_error_counter
        super().__init__(msg)

    async def on_consume(
        self,
        msg: "StreamMessage[Any]",
    ) -> "StreamMessage[Any]":
        self._received_messages_counter.labels(broker=self._broker_name).inc()
        return await super().on_consume(msg)

    async def after_processed(self, exc_type, exc_val, exc_tb):
        if exc_val is None:
            self._processed_messages_without_error_counter.labels(broker=self._broker_name).inc()
        else:
            self._processed_messages_with_error_counter.labels(broker=self._broker_name).inc()

        return await super().after_processed(exc_type, exc_val, exc_tb)