"""
Dramatiq Worker Starter Broker
===============================

Broker 初始化核心，提供 ReadableRedisBackend 和 init_broker 函数
"""

import logging

import dramatiq
from dramatiq.brokers.redis import RedisBroker
from dramatiq.middleware import AsyncIO, default_middleware
from dramatiq.results import Results
from dramatiq.results.backends import RedisBackend
from redis import Redis

from .middleware import TimingMiddleware, WorkerInfoMiddleware

logger = logging.getLogger(__name__)


class ReadableRedisBackend(RedisBackend):
    """
    可读的 Redis 结果后端。

    使用可读的结果键名格式：``{namespace}:{queue_name}:{actor_name}:{message_id}``
    """

    def build_message_key(self, message) -> str:
        """
        构建消息键名

        :param message: Dramatiq 消息对象
        :return: 格式化后的键名
        :rtype: str
        """
        return f"{self.namespace}:{message.queue_name}:{message.actor_name}:{message.message_id}"


def init_broker(
    redis_host: str = "localhost",
    redis_port: int = 6379,
    redis_db: int = 0,
    redis_password: str = "",
    redis_db_result: int = 1,
    namespace: str = "dramatiq-result",
    heartbeat_interval: int = 30,
    worker_ttl: int = 120,
    custom_middleware: list | None = None,
) -> RedisBroker:
    """
    初始化 Dramatiq Redis Broker

    :param redis_host: Redis 主机地址
    :type redis_host: str
    :param redis_port: Redis 端口
    :type redis_port: int
    :param redis_db: Redis 数据库编号
    :type redis_db: int
    :param redis_password: Redis 密码
    :type redis_password: str
    :param redis_db_result: 结果存储数据库编号
    :type redis_db_result: int
    :param namespace: 结果键命名空间
    :type namespace: str
    :param heartbeat_interval: Worker 心跳间隔（秒）
    :type heartbeat_interval: int
    :param worker_ttl: Worker 过期时间（秒）
    :type worker_ttl: int
    :param custom_middleware: 自定义中间件列表
    :type custom_middleware: list | None
    :return: RedisBroker 实例
    :rtype: RedisBroker
    """
    # 创建 Redis 客户端
    redis_client = Redis(
        host=redis_host,
        port=redis_port,
        db=redis_db,
        password=redis_password,
        decode_responses=False,
        encoding="utf-8",
        max_connections=10,
    )

    # 创建 Redis Broker
    redis_broker = RedisBroker(client=redis_client, middleware=[m() for m in default_middleware])

    # 创建可读的结果后端
    results_backend = ReadableRedisBackend(
        host=redis_host,
        port=redis_port,
        db=redis_db_result,
        password=redis_password,
        namespace=namespace,
    )

    # 添加默认中间件
    redis_broker.add_middleware(AsyncIO())

    # TimingMiddleware 必须在 Results 之前添加
    # after_process_message 执行顺序是后添加的先执行
    # 我们需要 TimingMiddleware 最后执行（以覆盖 Results 的结果）
    # 所以 TimingMiddleware 要先添加
    result_redis_client = Redis(
        host=redis_host,
        port=redis_port,
        db=redis_db_result,
        password=redis_password,
        decode_responses=False,
        encoding="utf-8",
        max_connections=10,
    )
    redis_broker.add_middleware(TimingMiddleware(redis_client=result_redis_client, namespace=namespace))

    # Results 后添加，这样它的 after_process_message 会先执行
    # 然后 TimingMiddleware 的 after_process_message 执行并覆盖结果
    redis_broker.add_middleware(Results(backend=results_backend))

    redis_broker.add_middleware(
        WorkerInfoMiddleware(redis_client=redis_client, heartbeat_interval=heartbeat_interval, worker_ttl=worker_ttl)
    )

    # 添加自定义中间件
    if custom_middleware:
        for middleware in custom_middleware:
            redis_broker.add_middleware(middleware)

    # 设置全局 Broker
    dramatiq.set_broker(redis_broker)

    logger.info(
        f"Broker initialized: redis://{redis_host}:{redis_port}/{redis_db} (results: {redis_db_result}, namespace: {namespace})"
    )

    return redis_broker
