import asyncio
import json
import os
import signal

import aio_pika
from aio_pika.exceptions import AMQPException

import metrics
from db import db_manager
from device_token_repository import database_lookup_by_user_id, database_lookup_by_uuid, remove_inactive_push_token
from secret_manager import cleanup_secret_files, return_credentials
from secret_manager.fernet_encryption import decrypt_token
from send_notification import send_notification
from simple_logger_handler import setup_logger

logger = setup_logger(__name__)


rmq_username = return_credentials("/etc/secrets/rmq_username")
rmq_password = return_credentials("/etc/secrets/rmq_password")
rmq_host = os.getenv("BACKEND_PN_RMQ_HOST", "localhost:30672")
rmq_vhost = os.getenv("BACKEND_PN_RMQ_VHOST", "app_notifications")
rmq_exchange = os.getenv("BACKEND_PN_RMQ_EXCHANGE", "app_notifications")

RABBITMQ_URL = f"amqp://{rmq_username}:{rmq_password}@{rmq_host}/{rmq_vhost}"


class RabbitMQConsumer:
	def __init__(self, url=RABBITMQ_URL, db_manager=db_manager, exchange_name=rmq_exchange,rmq_host=rmq_host):
		self.url = url
		self.rmq_host = rmq_host
		self.exchange_name = exchange_name
		self.connection: aio_pika.RobustConnection | None = None
		self.channel: aio_pika.RobustChannel | None = None
		self.exchange: aio_pika.Exchange | None = None
		self.queue: aio_pika.Queue | None = None
		self._closing = False
		self.db_manager = db_manager
		self.queue_name = "notifications"
		self.retry_queue_name = "notifications_retry"
		self.dlq_queue_name = "notifications_dlq"
		self.max_retries = 5
		self.retry_ttl = 5 * 60 * 10000
		logger.debug(f"[RabbitMQ] Initialized consumer for exchange {self.exchange_name}")


	async def connect(self):
		logger.info(f"[RabbitMQ] Connecting to {self.rmq_host}...")
		self.connection = await aio_pika.connect_robust(self.url)
		self.channel = await self.connection.channel()
		self.exchange = await self.channel.declare_exchange(
			self.exchange_name, aio_pika.ExchangeType.TOPIC, durable=True
		)
		self.queue = await self.channel.declare_queue(self.queue_name, durable=True)
		await self.queue.bind(self.exchange, routing_key="notify.user.*")
		logger.debug(f"[RabbitMQ] Bound queue '{self.queue_name}' to exchange '{self.exchange_name}'")

		retry_queue_args = {
				"x-message-ttl": self.retry_ttl,
				"x-dead-letter-exchange": self.exchange_name,
				"x-dead-letter-routing-key": "notify.user.retry",
			}
		logger.debug(f"[RabbitMQ] Declaring retry queue '{self.retry_queue_name}' with args: {retry_queue_args}")

		self.retry_queue = await self.channel.declare_queue(
			self.retry_queue_name, durable=True, arguments=retry_queue_args
		)
		logger.debug(f"[RabbitMQ] Retry queue '{self.retry_queue_name}' declared.")

		await self.channel.declare_queue(self.dlq_queue_name, durable=True)
		logger.debug(f"[RabbitMQ] DLQ queue '{self.dlq_queue_name}' declared.")

		metrics.start_metrics_server(port=9000)
		logger.debug("[Metrics] Prometheus metrics server started on port 9000.")

		self._metric_tasks = [
			asyncio.create_task(metrics.update_queue_gauge(self.channel, self.queue_name)),
			asyncio.create_task(metrics.update_queue_gauge(self.channel, self.retry_queue_name)),
			asyncio.create_task(metrics.update_queue_gauge(self.channel, self.dlq_queue_name)),
		]
		logger.debug(f"[Metrics] Gauge update tasks created for queues: {self.queue_name}, {self.retry_queue_name}, {self.dlq_queue_name}")

		logger.info("[RabbitMQ] Connected, exchange, retry, and DLQ queues ready.")

	async def send_message_to_dlq(self, uuid: str, message: aio_pika.IncomingMessage):
		data = json.loads(message.body.decode())
		data["uuid"] = uuid
		await self.channel.default_exchange.publish(
			aio_pika.Message(
				body=json.dumps(data).encode(),
				content_type="application/json",
				delivery_mode=aio_pika.DeliveryMode.PERSISTENT
			),
			routing_key=self.dlq_queue_name
		)
		logger.warning(f"[RabbitMQ] Message sent to DLQ: {data['uuid']}")
		metrics.MSG_FAILED.inc()

	async def send_message_to_retry_queue(self, uuid: str, message: aio_pika.IncomingMessage):
		"""
		Publishes a message to the retry queue with an incremented retry counter.

		Behaviour:
			- Each message is retried up to `max_retries` times.
			- After exceeding the limit, the message is forwarded to the DLQ.
			- The UUID of the intended device is embedded in the message payload.

		Args:
			uuid: Unique identifier for the device or token associated with the message.
			message: The original RabbitMQ message being retried.
		"""
		data = json.loads(message.body.decode())
		data["uuid"] = uuid
		retry_count = data.get("retry_count", 0) + 1
		data["retry_count"] = retry_count

		metrics.MSG_RETRY.labels(
			queue_name=self.retry_queue_name,
			uuid=uuid,
			retry_count=retry_count
		).inc()
		logger.debug(f"[RabbitMQ] Metrics updated for retry: uuid={uuid}, retry_count={retry_count}")

		if retry_count > self.max_retries:
			logger.debug(f"[RabbitMQ] Retry count exceeded max_retries={self.max_retries}, sending to DLQ")
			await self.send_message_to_dlq(uuid, message)
			return
		await self.channel.default_exchange.publish(
			aio_pika.Message(
			body=json.dumps(data).encode(),
			content_type="application/json",
			delivery_mode=aio_pika.DeliveryMode.PERSISTENT
			),
		routing_key=self.retry_queue_name
		)
		logger.info(f"[RabbitMQ] Message requeued for retry #{retry_count}")

	async def handle_message(self, message: aio_pika.IncomingMessage):
		"""
		Processes a single RabbitMQ message containing push notification data.

		Message routing:
			- Standard messages use the routing key: notify.user.<user_id>
			- Retry messages use the routing key: notify.user.retry and include a device UUID

		Behaviour:
			1. Parses and validates the message payload.
			2. Fetches encrypted push tokens from the database by UUID or user ID.
			3. Decrypts tokens and sends notifications via the Expo Push API.
			4. Delegates delivery result handling to `validate_delivery`.

		Error handling:
			- Invalid JSON → message discarded (not requeued)
			- RabbitMQ transport errors → message requeued
			- Unexpected exceptions → raised for shutdown or external handling
		"""
		if self._closing:
			logger.debug("[RabbitMQ] Skipping message because consumer is closing")
			return

		async with message.process():
			try:
				data = json.loads(message.body.decode())
				logger.debug(f"[RabbitMQ] Received message: {data}")

				if message.routing_key in [self.dlq_queue_name, "notify.dlq"]:
					logger.info(f"[RabbitMQ] Message in DLQ, skipping processing: {data}")
					return

				uuid = data.get("uuid")

				if uuid:
					encrypted_tokens = await database_lookup_by_uuid(uuid, db_manager)
				else:
					encrypted_tokens = await database_lookup_by_user_id(message.routing_key, db_manager)

				if not encrypted_tokens:
					logger.warning(f"[RabbitMQ] No push tokens found for routing key {message.routing_key}")
					return

				token_map = {row["uuid"]: row["token"].decode() for row in encrypted_tokens}
				for uuid, token  in token_map.items():
					decrypted_token = decrypt_token(token) 
					token_map[uuid] = decrypted_token
				response = await send_notification(message=data,push_tokens=token_map)

				# Validate delivery - any errors here should NOT prevent acknowledgment
				# because notifications were already sent. Instead, we handle retries/DLQ
				# by publishing new messages.
				try:
					await self.validate_delivery(response,message)
				except Exception as validation_error:
					logger.error(f"[RabbitMQ] Error during delivery validation: {validation_error}", exc_info=True)
					metrics.MSG_FAILED.inc()

			except json.JSONDecodeError as e:
				logger.error(f"[RabbitMQ] Bad message, discarding: {e}", exc_info=True)
				await message.nack(requeue=False)
			except AMQPException as e:
				logger.error(f"[RabbitMQ] RabbitMQ error: {e}", exc_info=True)
				await message.nack(requeue=True)
			except Exception as e:
				logger.critical(f"[RabbitMQ] Fatal error: {e}", exc_info=True)
				metrics.MSG_FAILED.inc()

	async def validate_delivery(self,response,message: aio_pika.IncomingMessage):
		"""
		Validates delivery results returned by the Expo Push API after a message is sent.

		This function inspects both the transport-level and API-level delivery statuses and
		takes appropriate follow-up actions:

			- Successful deliveries are logged and counted as published.
			- API-level errors trigger token expiration or redirection to the DLQ.
			- Failures (e.g., network or unknown issues) are requeued for retry.

		Right now the <push-receipt-id> is not stored for further validation but needs to be implemented if the userbase grows.

		Expected response structure:
			{
				"<uuid>": {
					"status": "ok" | "error" | "failure",
					"data": {
						"data": [
							{
								"status": "ok" | "error",
								"id": "<push-receipt-id>",
								"details": { "error": "<ErrorMessage>" } #Only exists if there is an error
							}
						]
					}
				}
			}

		See: https://docs.expo.dev/push-notifications/sending-notifications/
		"""
		for uuid, result in response.items():
			try:
				status = result.get("status")
				data_list = result.get("data", {}).get("data", [])

				api_status = None
				api_error = None
				if data_list and isinstance(data_list, list) and len(data_list) > 0:
					first_item = data_list[0]
					if isinstance(first_item, dict):
						api_status = first_item.get("status")
						api_error = first_item.get("details", {}).get("error")

				logger.debug(f"[RabbitMQ] Processing delivery for uuid={uuid}, status={status}, api_status={api_status}")

				if status == "ok" and api_status == "ok":
					logger.info(f"[RabbitMQ] Notification delivered successfully to {uuid}")
					metrics.MSG_PUBLISHED.inc()
					logger.debug(f"[RabbitMQ] Metrics updated for published message: uuid={uuid}")


				elif status == "ok" and api_status == "error":
					if api_error == "DeviceNotRegistered":
						expired = await remove_inactive_push_token(uuid, self.db_manager)
						if expired:
							logger.info(f"[RabbitMQ] Device no longer registered for uuid {uuid}, marked expired")
						else:
							logger.error(f"[RabbitMQ] Failed expiring token for uuid: {uuid}")
					else:
						await self.send_message_to_dlq(uuid, message)
						logger.debug(f"[RabbitMQ] Message sent to DLQ due to API error: uuid={uuid}, error={api_error}")


				elif status == "error":
					await self.send_message_to_dlq(uuid, message)
					logger.debug(f"[RabbitMQ] Message sent to DLQ due to status=error: uuid={uuid}")

				elif status == "failure":
					await self.send_message_to_retry_queue(uuid, message)
					logger.debug(f"[RabbitMQ] Message requeued for retry due to status=failure: uuid={uuid}")

			except Exception as e:
				logger.error(f"[RabbitMQ] Error processing delivery validation for uuid={uuid}: {e}", exc_info=True)
				try:
					await self.send_message_to_dlq(uuid,message)
				except Exception as dlq_error:
					logger.error(f"[RabbitMQ] Failed to send to DLQ for uuid={uuid}: {dlq_error}")


	async def consume(self):
		if self._closing:
			logger.debug("[RabbitMQ] Consumer is closing, not starting consume")
			return
		if not self.queue:
			raise RuntimeError("[RabbitMQ] Queue not initialized")
		await self.queue.consume(self.handle_message, no_ack=False)
		logger.info("[RabbitMQ] Consuming messages...")

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



async def main():
	logger.info("[APP] Starting application")
	await db_manager.connect()
	consumer = RabbitMQConsumer(db_manager=db_manager)
	await consumer.connect()
	await consumer.consume()
	stop_event = asyncio.Event()

	SECRET_PATHS = frozenset({
		"/etc/secrets/encryption_key",
		"/etc/secrets/db_username",
		"/etc/secrets/db_password",
		"/etc/secrets/rmq_username",
		"/etc/secrets/rmq_password"
	})
	cleanup_secret_files(SECRET_PATHS)

	for sig in (signal.SIGINT, signal.SIGTERM):
		asyncio.get_running_loop().add_signal_handler(sig, stop_event.set)

	await stop_event.wait()
	logger.info("[APP] Stopping application")
	await consumer.close()
	await db_manager.close()

if __name__ == "__main__":
	asyncio.run(main())