import asyncio
import aio_pika
from aio_pika.exceptions import AMQPException
from secret_manager import return_credentials
import os
from simple_logger_handler import setup_logger
import json
import signal
from urllib.parse import quote
import weakref

logger = setup_logger(__name__)

class RabbitMQConsumer:
	"""Pull-based RabbitMQ consumer. Messages are fetched on demand, one queue per user."""
	def __init__(self):
		self.rmq_host = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_HOST", "localhost")
		self.rmq_vhost = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_VHOST", "app_notifications")
		self.rmq_exchange = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_EXCHANGE", "app_notifications")
		self.pull_queue_prefix = os.getenv("BACKEND_PULL_NOTIFICATIONS_RMQ_QUEUE_PREFIX", "notifications-pull")
		self.pull_routing_prefix = os.getenv("BACKEND_PULL_NOTIFICATIONS_RMQ_ROUTING_PREFIX", "notify.pull")

		self._user_locks: weakref.WeakValueDictionary[int, asyncio.Lock] = weakref.WeakValueDictionary()
		self.connection: aio_pika.RobustConnection | None = None
		self.channel: aio_pika.RobustChannel | None = None
		self.exchange: aio_pika.Exchange | None = None
		self._closing = False
		logger.debug(f"[RabbitMQ] Initialized consumer for exchange {self.rmq_exchange}")

	def _lock_for(self, user_id: int) -> asyncio.Lock:
		"""
		Return the asyncio lock for this user, creating it if needed. Locks are held
		weakly, so they are cleaned up automatically once nobody is using them.
		"""
		lock = self._user_locks.get(user_id)
		if lock is None:
			lock = asyncio.Lock()
			self._user_locks[user_id] = lock
		return lock

	async def connect(self):
		username = return_credentials("/etc/secrets/rmq_username")
		password = return_credentials("/etc/secrets/rmq_password")
		url = (
			f"amqp://{quote(username, safe='')}:{quote(password, safe='')}"
			f"@{self.rmq_host}/{quote(self.rmq_vhost, safe='')}"
		)

		logger.info(f"[RabbitMQ] Connecting to {self.rmq_host}...")
		self.connection = await aio_pika.connect_robust(url)
		self.channel = await self.connection.channel()
		self.exchange = await self.channel.declare_exchange(
			self.rmq_exchange, aio_pika.ExchangeType.TOPIC, durable=True
		)
		logger.info("[RabbitMQ] Connected to RabbitMQ")

	async def ensure_user_queue(self, user_id: int) -> aio_pika.Queue:
		"""
		Declare the user's durable queue (with expiry and message TTL) and bind it to
		the exchange with the user's routing key. Safe to call repeatedly, recreates it if it expired, and resets its
		idle-expiry timer.
		"""
		queue = await self.channel.declare_queue(
			f"{self.pull_queue_prefix}.{user_id}",
			durable=True,
			arguments={
				"x-expires": 30 * 24 * 3600 * 1000,
				"x-message-ttl": 7 * 24 * 3600 * 1000,
			},
		)
		await queue.bind(self.exchange, routing_key=f"{self.pull_routing_prefix}.{user_id}")
		return queue
		
	async def fetch_messages(self, user_id: int, max_messages: int = 100) -> list[dict]:
		"""Drain up to max_messages queued messages for this user and return their
		decoded JSON bodies.

		- Only one fetch per user runs at a time; different users run concurrently.
		- Messages are acked after the whole batch has been read, so an error mid-read
		  puts the batch back on the queue and re-raises.
		- Undecodable messages are rejected without requeue (discarded).
		- Returns an empty list if nothing is queued.
		- Delivery is at-most-once towards the client: messages are acked before the
		  HTTP response is sent.
		"""

		if self._closing:
			raise RuntimeError("[RabbitMQ] Consumer is closing")
		if not self.channel:
			raise RuntimeError("[RabbitMQ] Queue not initialized, call connect() first")

		user_id = int(user_id)

		async with self._lock_for(user_id):  #one drain at a time
			queue = await self.ensure_user_queue(user_id)
			results: list[dict] = []
			to_ack: list[aio_pika.IncomingMessage] = []
			try:
				while len(to_ack) < max_messages:
					message = await queue.get(no_ack=False, fail=False, timeout=5)
					if message is None:
						break

					try:
						data = json.loads(message.body.decode())
					except (json.JSONDecodeError, UnicodeDecodeError) as e:
						logger.error(f"[RabbitMQ] Bad message, discarding: {e}", exc_info=True)
						await message.reject(requeue=False)
						continue

					results.append(data)
					to_ack.append(message)

				for message in to_ack:
					await message.ack()

				logger.debug(f"[RabbitMQ] Fetched {len(results)} message(s)")
				return results

			except (AMQPException, OSError) as e:
				logger.error(f"[RabbitMQ] Error while fetching, requeueing batch: {e}", exc_info=True)
				for message in to_ack:
					try:
						await message.nack(requeue=True)
					except Exception:
						pass  # channel is gone; the broker requeues unacked messages itself
				raise

	async def close(self):
		self._closing = True
		if self.connection:
			await self.connection.close()
			logger.info("[RabbitMQ] Connection closed")