# coding=utf-8
"""
@author: noybzy
@time: 2025/5/13 14:34
@file: mq.py
@describe: RabbitMQ 客户端封装模块
"""
import functools
import threading
import time
from queue import Empty, Queue
from typing import Callable, Optional, TYPE_CHECKING

from loguru import logger
from nosp.config import RabbitMQConfig
from nosp.lazy import LazyLoader

if TYPE_CHECKING:
    import pika
else:
    pika = LazyLoader("pika")


class Message:
    """消息封装类，提供 ack/nack/reject 方法"""

    def __init__(self, ch, delivery_tag: int, body: bytes):
        self.body = body
        self.delivery_tag = delivery_tag
        self.ch = ch

    def __ack_message(self, ch, delivery_tag: int):
        if ch and ch.is_open:
            ch.basic_ack(delivery_tag)

    def __nack_message(self, ch, delivery_tag: int, requeue: bool):
        if ch and ch.is_open:
            ch.basic_nack(delivery_tag, requeue=requeue)

    def __reject_message(self, ch, delivery_tag: int, requeue: bool):
        if ch and ch.is_open:
            ch.basic_reject(delivery_tag, requeue=requeue)

    def ack(self):
        """确认消息"""
        cb = functools.partial(self.__ack_message, self.ch, self.delivery_tag)
        self.ch.connection.add_callback_threadsafe(cb)

    def nack(self, requeue: bool = True):
        """拒绝消息，可重新入队"""
        try:
            cb = functools.partial(self.__nack_message, self.ch, self.delivery_tag, requeue)
            self.ch.connection.add_callback_threadsafe(cb)
        except Exception as e:
            logger.error(f"nack 失败：{e}")

    def reject(self, requeue: bool = True):
        """拒绝消息，可重新入队"""
        try:
            cb = functools.partial(self.__reject_message, self.ch, self.delivery_tag, requeue)
            self.ch.connection.add_callback_threadsafe(cb)
        except Exception as e:
            logger.error(f"reject 失败：{e}")


class RabbitMQ:
    """RabbitMQ 客户端，支持消息生产和消费"""

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 5672,
        username: str = "admin",
        password: str = "123456",
        queue_name: str = "test",
        virtual_host: str = "/",
    ):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.queue_name = queue_name
        self.virtual_host = virtual_host

        self._lock = threading.Lock()
        self._connection = None
        self._channel = None
        self._work_queue: Queue = Queue()
        self._callback: Optional[Callable[[Message], None]] = None
        self._running = False

    @staticmethod
    def simple(conf: RabbitMQConfig, queue_name: str = "test") -> "RabbitMQ":
        """通过配置对象快速创建实例"""
        return RabbitMQ(
            conf.host,
            conf.port,
            conf.username,
            conf.password,
            queue_name,
            conf.virtual_host,
        )

    def create_queue(self, queue_name: Optional[str] = None, dlq_name: Optional[str] = None) -> None:
        """
        创建队列，支持死信队列配置

        :param queue_name: 队列名称，默认为实例配置的 queue_name
        :param dlq_name: 死信队列名称，可选
        """
        queue_name = queue_name or self.queue_name
        self.ensure_channel()

        arguments = None
        if dlq_name:
            # 声明死信队列
            self._channel.queue_declare(queue=dlq_name, durable=True)
            arguments = {
                "x-dead-letter-exchange": "",  # 使用默认交换机
                "x-dead-letter-routing-key": dlq_name,  # 死信消息发送到死信队列的路由键
            }

        # 声明主队列
        self._channel.queue_declare(
            queue=queue_name,
            durable=True,
            arguments=arguments,
        )
        logger.info(f"队列 '{queue_name}' 创建成功" + (f"，死信队列 '{dlq_name}'" if dlq_name else ""))

    def _create_connection(self) -> "pika.BlockingConnection":
        """创建一个新的连接"""
        credentials = pika.PlainCredentials(self.username, self.password)
        parameters = pika.ConnectionParameters(
            host=self.host,
            port=self.port,
            virtual_host=self.virtual_host,
            credentials=credentials,
            heartbeat=600,
            blocked_connection_timeout=300,
        )
        return pika.BlockingConnection(parameters)

    def ensure_channel(self) -> None:
        """确保 channel 可用，若不可用则重新建立连接和 channel"""
        with self._lock:
            if not self._connection or self._connection.is_closed:
                self._connection = self._create_connection()
                self._channel = self._connection.channel()
            elif not self._channel or self._channel.is_closed:
                self._channel = self._connection.channel()

    def _safe_close(self) -> None:
        """安全地关闭连接和 channel"""
        with self._lock:
            if self._channel and self._channel.is_open:
                try:
                    self._channel.close()
                except pika.exceptions.ConnectionWrongStateError:
                    pass
            self._channel = None

            if self._connection and not self._connection.is_closed:
                try:
                    self._connection.close()
                except pika.exceptions.ConnectionWrongStateError:
                    pass
            self._connection = None

    def close(self) -> None:
        """手动关闭连接，用于优雅退出"""
        self._running = False
        self._safe_close()
        logger.info("RabbitMQ 连接已关闭")

    def consume(self, callback: Callable[[Message], None], prefetch_count: int = 1) -> None:
        """
        启动消费循环

        :param callback: 消息处理回调函数
        :param prefetch_count: 预取消息数量
        """
        self._callback = callback
        self._running = True

        while self._running:
            try:
                self.ensure_channel()
                self._channel.basic_qos(prefetch_count=prefetch_count)
                self._channel.basic_consume(
                    on_message_callback=self._on_message,
                    queue=self.queue_name,
                    auto_ack=False,
                )
                logger.info(f"开始消费队列 '{self.queue_name}'，预取数：{prefetch_count}")
                self._channel.start_consuming()
            except (pika.exceptions.AMQPConnectionError, pika.exceptions.ConnectionClosedByBroker) as e:
                self._safe_close()
                logger.warning(f"连接错误：{e}，将在 5 秒后重试...")
                time.sleep(5)
            except KeyboardInterrupt:
                self._safe_close()
                logger.warning("手动终止消费")
                break
            except Exception as e:
                self._safe_close()
                logger.error(f"未知错误：{e}，将在 5 秒后重试...")
                time.sleep(5)

    def produce(self, body: str, queue_name: Optional[str] = None) -> bool:
        """
        向队列发送消息

        :param body: 消息体
        :param queue_name: 队列名称，默认为实例配置的 queue_name
        :return: 发送是否成功
        """
        queue_name = queue_name or self.queue_name

        while True:
            try:
                self.ensure_channel()
                self._channel.basic_publish(
                    exchange="",
                    routing_key=queue_name,
                    body=str(body),
                    properties=pika.BasicProperties(delivery_mode=2),  # 持久化消息
                )
                logger.debug(f"消息已发送到队列 '{queue_name}'")
                return True
            except (pika.exceptions.AMQPConnectionError, pika.exceptions.ConnectionClosedByBroker) as e:
                self._safe_close()
                logger.warning(f"生产者连接错误：{e}，将在 5 秒后重试...")
                time.sleep(5)
            except Exception as e:
                self._safe_close()
                logger.error(f"发送失败：{e}，将在 5 秒后重试...")
                time.sleep(5)

    def _on_message(self, ch, method_frame, _header_frame, body: bytes) -> None:
        """消息到达回调，将消息放入工作队列"""
        delivery_tag = method_frame.delivery_tag
        self._work_queue.put((ch, delivery_tag, body))

    def _process_message(self, ch, delivery_tag: int, body: bytes) -> None:
        """处理单条消息"""
        try:
            message = Message(ch, delivery_tag, body)
            if self._callback:
                self._callback(message)
        except Exception as e:
            logger.error(f"消息处理失败：{e}")
            # 处理失败时 nack 消息，允许重新入队
            try:
                if ch and ch.is_open:
                    ch.basic_nack(delivery_tag, requeue=True)
            except Exception as nack_error:
                logger.error(f"nack 失败：{nack_error}")

    def _work_loop(self) -> None:
        """工作线程循环，从工作队列获取并处理消息"""
        while self._running:
            try:
                ch, delivery_tag, body = self._work_queue.get(timeout=1.0)
                self._process_message(ch, delivery_tag, body)
                self._work_queue.task_done()
            except Empty:
                continue
            except Exception as e:
                logger.error(f"工作线程异常：{e}")

    def start_consuming(self, callback: Callable[[Message], None], num_workers: int = 1) -> None:
        """
        启动多线程消费

        :param callback: 消息处理回调函数
        :param num_workers: 工作线程数量
        """
        self._callback = callback
        self._running = True

        # 启动工作线程池
        for i in range(num_workers):
            t = threading.Thread(target=self._work_loop, name=f"RabbitMQ-Worker-{i}", daemon=True)
            t.start()
            logger.debug(f"工作线程 {i} 已启动")

        # 启动消费循环（阻塞）
        self.consume(callback, prefetch_count=num_workers + 1)
