"""RabbitMQ as an on-prem queue, through ``pika`` (add ``pika`` to requirements.txt).

Template: not run in CI. Delivery is at-least-once. A message is acked only
after the handler succeeded; a failure is nacked back to the queue, or
published to the dead-letter queue once it has failed ``max_receive`` times.

Use it with::

    resources:
      jobs:
        type: queue
        external: true
        options:
          driver: custom
          class: pskit_drivers.rabbitmq_queue:RabbitQueue
          url_env: RABBITMQ_URL          # amqp://user:pass@rabbit:5672/%2F
          queue: jobs

RabbitMQ has no visibility timeout: an unacked message stays reserved until it
is acked/nacked or the connection drops (then RabbitMQ redelivers it). Receive
counts come from ``x-delivery-count``, which quorum queues set. On classic queues
a redelivery only counts as one retry.
"""

from __future__ import annotations

import threading
import time
import uuid
from typing import Any, Dict, List, Optional, Set


class RabbitQueue:
    def __init__(self, url: str, queue: str, durable: bool = True, prefetch: int = 50) -> None:
        import pika  # imported here so pskit can package the app without pika installed

        self._pika = pika
        self.params = pika.URLParameters(url)
        self.queue = queue
        self.durable = durable
        self.prefetch = prefetch
        self._consume: Any = None  # used only by the host's poller thread (claim/settle)
        self._publish: Any = None  # submit() may come from HTTP threads
        self._publish_lock = threading.Lock()

    def _channel(self, which: str) -> Any:
        current = getattr(self, which)
        if current is not None and current.is_open:
            return current
        channel = self._pika.BlockingConnection(self.params).channel()
        channel.queue_declare(queue=self.queue, durable=self.durable)
        if which == "_consume":
            channel.basic_qos(prefetch_count=self.prefetch)
        setattr(self, which, channel)
        return channel

    def submit(self, body: str, attributes: Dict[str, Any]) -> str:
        ident = uuid.uuid4().hex
        props = self._pika.BasicProperties(
            message_id=ident,
            delivery_mode=2 if self.durable else 1,
            headers={str(k): str(v) for k, v in (attributes or {}).items()},
            timestamp=int(time.time()),
        )
        with self._publish_lock:
            self._channel("_publish").basic_publish(exchange="", routing_key=self.queue, body=body, properties=props)
        return ident

    def claim(self, batch: int, visibility: float, now: float) -> List[Dict[str, Any]]:
        channel = self._channel("_consume")
        claimed: List[Dict[str, Any]] = []
        for _ in range(max(1, batch)):
            method, props, body = channel.basic_get(queue=self.queue, auto_ack=False)
            if method is None:
                break
            headers = dict(props.headers or {})
            receives = int(headers.pop("x-delivery-count", 1 if method.redelivered else 0))
            claimed.append(
                {
                    "id": props.message_id or str(method.delivery_tag),
                    "body": body.decode("utf-8", "replace"),
                    "attributes": {k: str(v) for k, v in headers.items()},
                    "receives": receives,
                    "enqueued_at": float(props.timestamp or now),
                    "_tag": method.delivery_tag,
                }
            )
        return claimed

    def settle(
        self,
        claimed: List[Dict[str, Any]],
        failed_ids: Set[str],
        max_receive: int,
        dead_letter: Optional[Any],
        now: float,
    ) -> None:
        channel = self._channel("_consume")
        for msg in claimed:
            if msg["id"] not in failed_ids:
                channel.basic_ack(delivery_tag=msg["_tag"])
            elif dead_letter is not None and int(msg["receives"]) + 1 >= max_receive:
                dead_letter.submit(msg["body"], dict(msg["attributes"], pskit_source_queue=self.queue))
                channel.basic_ack(delivery_tag=msg["_tag"])  # only after the DLQ has it
            else:
                channel.basic_nack(delivery_tag=msg["_tag"], requeue=True)

    def depth(self) -> int:
        with self._publish_lock:
            frame = self._channel("_publish").queue_declare(queue=self.queue, durable=self.durable, passive=True)
        return int(frame.method.message_count)

    def recover(self, now: float, visibility: float) -> None:
        return None  # RabbitMQ redelivers unacked messages when a connection drops
