# rabbit_base.py
from typing import Callable, Dict, List, Optional
import aio_pika
from urllib.parse import urlsplit, unquote
from config import settings

ExchangeResolver = Callable[[str], str]  # exchangeName -> exchangeType


def _parse_amqp_url(url: str) -> dict:
    """
    Convert an AMQP URL into kwargs for aio_pika.connect_robust,
    so we don't pass the raw URL (and leak the password in logs).
    """
    parts = urlsplit(url)
    return {
        "host": parts.hostname or "localhost",
        "port": parts.port or (5671 if parts.scheme == "amqps" else 5672),
        "login": parts.username or "guest",
        "password": parts.password or "guest",
        "virtualhost": unquote(parts.path[1:] or "/"),
        "ssl": parts.scheme == "amqps",
    }


class RabbitBase:
    def __init__(self, exchange_type_resolver: Optional[ExchangeResolver] = None):
        self._conn: Optional[aio_pika.RobustConnection] = None
        self._chan: Optional[aio_pika.RobustChannel] = None
        self._exchanges: Dict[str, aio_pika.Exchange] = {}
        self._exchange_type_resolver = exchange_type_resolver or (
            lambda _: settings.RABBIT_EXCHANGE_TYPE
        )

    async def connect(self) -> None:
        if self._conn and not self._conn.is_closed:
            return
        conn_kwargs = _parse_amqp_url(str(settings.AMQP_URL))
        self._conn = await aio_pika.connect_robust(**conn_kwargs)
        self._chan = await self._conn.channel()
        await self._chan.set_qos(prefetch_count=settings.RABBIT_PREFETCH)

    async def ensure_exchange(self, name: str) -> aio_pika.Exchange:
        await self.connect()
        if name in self._exchanges:
            return self._exchanges[name]
        ex_type = self._exchange_type_resolver(name)
        ex = await self._chan.declare_exchange(
            name, getattr(aio_pika.ExchangeType, ex_type), durable=True
        )
        self._exchanges[name] = ex
        return ex

    async def declare_queue_bind(
        self,
        exchange: str,
        queue_name: str,
        routing_keys: List[str],
        ttl_ms: Optional[int],
    ):
        await self.connect()
        ex = await self.ensure_exchange(exchange)
        args: Dict[str, int] = {}
        if ttl_ms:
            args["x-message-ttl"] = ttl_ms
        q = await self._chan.declare_queue(
            queue_name,
            durable=True,
            exclusive=False,
            auto_delete=True,
            arguments=args,
        )
        for rk in routing_keys or [""]:
            await q.bind(ex, rk)
        return q