from typing import Type, Protocol

from faststream._internal.broker import BrokerUsecase
from faststream.asgi import AsgiFastStream
from faststream.prometheus import PrometheusMiddleware
from prometheus_client import make_asgi_app

from ..core import MetricsExporter


class MiddlewareFactory(Protocol):
    def __call__(self, metrics: MetricsExporter) -> PrometheusMiddleware: ...


_BROKER_MIDDLEWARE_MAP: dict[Type[BrokerUsecase], MiddlewareFactory] = {}

# --- RabbitMQ ---
try:
    from faststream.rabbit import RabbitBroker
    from faststream.rabbit.prometheus import RabbitPrometheusMiddleware

    class CustomRabbitMiddleware(RabbitPrometheusMiddleware):
        def __init__(self, metrics: MetricsExporter) -> None:
            super().__init__(registry=metrics.get_registry(), metrics_prefix=metrics.prefix)

    _BROKER_MIDDLEWARE_MAP[RabbitBroker] = CustomRabbitMiddleware
except ImportError:
    pass

# --- Kafka ---
try:
    from faststream.kafka import KafkaBroker
    from faststream.kafka.prometheus import KafkaPrometheusMiddleware

    class CustomKafkaMiddleware(KafkaPrometheusMiddleware):
        def __init__(self, metrics: MetricsExporter) -> None:
            super().__init__(registry=metrics.get_registry(), metrics_prefix=metrics.prefix)

    _BROKER_MIDDLEWARE_MAP[KafkaBroker] = CustomKafkaMiddleware
except ImportError:
    pass

# --- Redis ---
try:
    from faststream.redis import RedisBroker
    from faststream.redis.prometheus import RedisPrometheusMiddleware

    class CustomRedisMiddleware(RedisPrometheusMiddleware):
        def __init__(self, metrics: MetricsExporter) -> None:
            super().__init__(registry=metrics.get_registry(), metrics_prefix=metrics.prefix)

    _BROKER_MIDDLEWARE_MAP[RedisBroker] = CustomRedisMiddleware
except ImportError:
    pass

try:
    from faststream.confluent import KafkaBroker as ConfluentBroker
    from faststream.confluent.prometheus import KafkaPrometheusMiddleware as ConfluentPrometheusMiddleware

    class CustomConfluentPrometheusMiddleware(ConfluentPrometheusMiddleware):
        def __init__(self, metrics: MetricsExporter) -> None:
            super().__init__(registry=metrics.get_registry(), metrics_prefix=metrics.prefix)

    _BROKER_MIDDLEWARE_MAP[ConfluentBroker] = CustomConfluentPrometheusMiddleware
except ImportError:
    pass

# --- Nats ---
try:
    from faststream.nats import NatsBroker
    from faststream.nats.prometheus import NatsPrometheusMiddleware

    class CustomNatsPrometheusMiddleware(NatsPrometheusMiddleware):
        def __init__(self, metrics: MetricsExporter):
            super().__init__(registry=metrics.get_registry(), metrics_prefix=metrics.prefix)

    _BROKER_MIDDLEWARE_MAP[NatsBroker] = CustomNatsPrometheusMiddleware

except ImportError:
    pass


def setup_broker_metrics(broker: BrokerUsecase, metrics: MetricsExporter) -> None:
    broker_type = type(broker)
    middleware_cls = _BROKER_MIDDLEWARE_MAP.get(broker_type)

    if not middleware_cls:
        for b_cls, m_cls in _BROKER_MIDDLEWARE_MAP.items():
            if isinstance(broker, b_cls):
                middleware_cls = m_cls
                break

    if not middleware_cls:
        available = [k.__name__ for k in _BROKER_MIDDLEWARE_MAP.keys()]
        raise ValueError(
            f"Unsupported broker type: {broker_type.__name__}. " f"Installed drivers support: {', '.join(available)}"
        )
    broker.add_middleware(middleware_cls(metrics=metrics))


def setup_faststream_metrics(app: AsgiFastStream, metrics: MetricsExporter) -> None:
    broker: BrokerUsecase = app.broker
    setup_broker_metrics(broker, metrics)
    app.routes.append(("/metrics", make_asgi_app(metrics.get_registry())))
