"""
Dramatiq Worker Starter Middleware
==================================

包含 TimingMiddleware 和 WorkerInfoMiddleware
"""

import json
import logging
import os
import time
import uuid
from datetime import datetime
from threading import Thread
from typing import Any

import dramatiq
from dramatiq.middleware import Middleware
from redis import Redis

from .models import ErrorInfo, ResultData, TaskMetadata

logger = logging.getLogger(__name__)


class TimingMiddleware(Middleware):
    """
    任务执行时间记录中间件。

    功能：

    * 在任务开始执行时记录开始时间
    * 在任务完成时记录结束时间和持续时间
    * 直接更新 Redis 中存储的结果，添加时间信息

    数据结构 ``_timing``::

        {
            'start_datetime': str,  # 开始时间（ISO 格式）
            'end_datetime': str,    # 结束时间（ISO 格式）
            'duration_ms': int,     # 持续时间（毫秒）
            'queue_name': str,      # 队列名称
            'actor_name': str,      # Actor 名称
            'worker_id': str,       # Worker ID
            'exception': str,       # 异常信息（如果有）
        }
    """

    def __init__(self, redis_client: Redis | None = None, namespace: str = "dramatiq-result"):
        """
        初始化时间记录中间件

        :param redis_client: Redis 客户端（连接到结果数据库）
        :type redis_client: Redis | None
        :param namespace: 结果键命名空间
        :type namespace: str
        """
        self.runtime_data = {}
        self.worker_id = str(uuid.uuid4())[:8]
        self.redis_client = redis_client
        self.namespace = namespace

    def before_process_message(self, broker: dramatiq.Broker, message: dramatiq.Message) -> None:
        """
        任务执行前回调 - 记录开始时间

        :param broker: Dramatiq Broker
        :type broker: dramatiq.Broker
        :param message: 任务消息
        :type message: dramatiq.Message
        """
        start_datetime = datetime.now()
        self.runtime_data[message.message_id] = {
            "start_datetime": start_datetime,
            "start_timestamp_ms": int(time.time() * 1000),
            "queue_name": message.queue_name,
            "actor_name": message.actor_name,
            "worker_id": self.worker_id,
        }

    def after_process_message(
        self,
        broker: dramatiq.Broker,
        message: dramatiq.MessageProxy,
        *,
        result: Any | None = None,
        exception: BaseException | None = None,
    ) -> None:
        """
        任务执行后回调 - 记录结束时间并更新 Redis 中的结果。

        :param broker: Dramatiq Broker
        :type broker: dramatiq.Broker
        :param message: 任务消息
        :type message: dramatiq.MessageProxy
        :param result: 执行结果
        :type result: Any | None
        :param exception: 异常信息（如果有）
        :type exception: BaseException | None
        """
        if message.message_id not in self.runtime_data:
            return

        end_datetime = datetime.now()
        end_timestamp_ms = int(time.time() * 1000)
        runtime_info = self.runtime_data[message.message_id]

        # 计算持续时间（毫秒）
        duration = end_timestamp_ms - runtime_info["start_timestamp_ms"]

        # 记录结束时间和持续时间
        runtime_info["end_datetime"] = end_datetime
        runtime_info["duration_ms"] = duration
        runtime_info["exception"] = str(exception) if exception else None

        # 直接更新 Redis 中存储的结果（成功和失败都处理）
        if self.redis_client:
            try:
                # 构建结果键
                result_key = f"dramatiq-result:{message.queue_name}:{message.actor_name}:{message.message_id}"

                # 使用 Pydantic 模型构建元数据
                task_metadata = TaskMetadata.from_runtime_info(
                    start_datetime=runtime_info["start_datetime"],
                    end_datetime=end_datetime,
                    duration_ms=duration,
                    queue_name=runtime_info["queue_name"],
                    actor_name=runtime_info["actor_name"],
                    worker_id=runtime_info["worker_id"],
                    exception=runtime_info["exception"],
                )

                # 构建统一格式的结果数据
                if exception:
                    # 失败时，构建包含错误信息的结果
                    result_data = ResultData(
                        result=None,
                        error=ErrorInfo(type=type(exception).__name__, message=str(exception)),
                        metadata=task_metadata,
                    )
                    logger.info(f"result_data: {result_data}")
                    logger.info(f"result: {result}")
                else:
                    # 成功时，包装结果
                    result_data = ResultData(
                        result=result,
                        metadata=task_metadata,
                    )

                # 使用 lpush 存储到列表头部（确保我们的格式最先被读取）
                # 然后删除其他条目，避免重复
                logger.info(f"Storing result with timing: {result_key=}")
                self.redis_client.lpush(result_key, result_data.model_dump_json())
                # 只保留最新的条目
                self.redis_client.ltrim(result_key, 0, 0)
            except Exception as e:
                logger.error(f"[Timing] Failed to update result in Redis: {e}", exc_info=True)

        # 清理临时数据
        # 如果有异常，不清理（让 after_nack 处理）
        # 如果成功，清理
        if exception is None and message.message_id in self.runtime_data:
            del self.runtime_data[message.message_id]

    def after_nack(self, broker: dramatiq.Broker, message: dramatiq.MessageProxy) -> None:
        """
        消息被拒绝后回调 - 处理失败情况。

        Dramatiq 的 Results 中间件在任务失败时调用 ``after_nack`` 而不是 ``after_process_message``。
        我们需要在这里覆盖 Results 存储的错误格式。

        :param broker: Dramatiq Broker
        :type broker: dramatiq.Broker
        :param message: 任务消息
        :type message: dramatiq.MessageProxy
        """
        if message.message_id not in self.runtime_data:
            return

        end_datetime = datetime.now()
        end_timestamp_ms = int(time.time() * 1000)
        runtime_info = self.runtime_data[message.message_id]

        # 计算持续时间（毫秒）
        duration = end_timestamp_ms - runtime_info["start_timestamp_ms"]

        # 获取异常信息
        exception = getattr(message, "_exception", None) or message.options.get("exception")
        exception_str = str(exception) if exception else "Unknown error"
        exception_type = type(exception).__name__ if exception else "UnknownError"

        # 记录结束时间和持续时间
        runtime_info["end_datetime"] = end_datetime
        runtime_info["duration_ms"] = duration
        runtime_info["exception"] = exception_str

        # 更新 Redis 中的结果（覆盖 Dramatiq 的错误格式）
        if self.redis_client:
            try:
                # 构建结果键
                result_key = f"dramatiq-result:{message.queue_name}:{message.actor_name}:{message.message_id}"

                # 使用 Pydantic 模型构建元数据
                task_metadata = TaskMetadata.from_runtime_info(
                    start_datetime=runtime_info["start_datetime"],
                    end_datetime=end_datetime,
                    duration_ms=duration,
                    queue_name=runtime_info["queue_name"],
                    actor_name=runtime_info["actor_name"],
                    worker_id=runtime_info["worker_id"],
                    exception=exception_str,
                )

                # 构建统一格式的结果数据
                result_data = ResultData(
                    result=None,
                    error=ErrorInfo(type=exception_type, message=exception_str),
                    metadata=task_metadata,
                )

                logger.info(f"[Timing] after_nack storing result with timing: {result_key=}")
                # 使用 lpush 存储到列表头部
                self.redis_client.lpush(result_key, result_data.model_dump_json())
                # 只保留最新的条目
                self.redis_client.ltrim(result_key, 0, 0)
            except Exception as e:
                logger.error(f"[Timing] Failed to update result in after_nack: {e}", exc_info=True)

        # 清理临时数据
        del self.runtime_data[message.message_id]


class WorkerInfoMiddleware(Middleware):
    """
    Worker 信息注册中间件。

    功能：

    * Worker 启动时自动注册到 Redis
    * 定期发送心跳更新 Worker 状态
    * Worker 停止时自动注销
    * 将 Worker 支持的 Actor 列表注册到 Redis

    Redis 数据结构：

    * ``dramatiq:workers:active`` (ZSET): ``{worker_id: timestamp}``
    * ``dramatiq:worker:{worker_id}:actors`` (SET): ``{actor_name, ...}``
    * ``dramatiq:worker:{worker_id}:info`` (HASH)::

          - id: worker_id
          - hostname: 主机名
          - pid: 进程ID
          - started_at: 启动时间戳
          - last_heartbeat: 最后心跳时间戳
          - actors: Actor 列表 (JSON)
    """

    def __init__(
        self,
        redis_client: Redis | None = None,
        heartbeat_interval: int = 30,
        worker_ttl: int = 120,
    ):
        """
        初始化中间件

        :param redis_client: Redis 客户端（从 broker 获取）
        :type redis_client: Redis | None
        :param heartbeat_interval: 心跳间隔（秒）
        :type heartbeat_interval: int
        :param worker_ttl: Worker 过期时间（秒），应为 heartbeat_interval 的 4 倍
        :type worker_ttl: int
        """
        self.redis_client = redis_client
        self.heartbeat_interval = heartbeat_interval
        self.worker_ttl = worker_ttl
        self.worker_id = str(uuid.uuid4())
        self._heartbeat_thread: Thread | None = None
        self._running = False

        # Worker 元信息
        self.hostname = os.uname().nodename
        self.pid = os.getpid()
        self.started_at = int(time.time() * 1000)

        # Actor 列表（延迟获取）
        self.actors = []

    def after_worker_boot(self, broker: dramatiq.Broker, worker: dramatiq.Worker) -> None:
        """
        Worker 启动后回调

        :param broker: Dramatiq Broker
        :type broker: dramatiq.Broker
        :param worker: Dramatiq Worker
        :type worker: dramatiq.Worker
        """
        # 从 broker 获取 Redis 客户端（如果未提供）
        if self.redis_client is None:
            self.redis_client = broker.client

        # 获取 Worker 支持的 Actor 列表
        self.actors = list(broker.actors.keys())

        # 注册 Worker
        self._register_worker()

        # 启动心跳线程
        self._running = True
        self._heartbeat_thread = Thread(
            target=self._heartbeat_loop, daemon=True, name=f"worker-heartbeat-{self.worker_id[:8]}"
        )
        self._heartbeat_thread.start()

        logger.info(f"Worker {self.worker_id[:8]} registered with {len(self.actors)} actors: {self.actors}")

    def before_worker_shutdown(self, broker: dramatiq.Broker, worker: dramatiq.Worker) -> None:
        """
        Worker 停止前回调

        :param broker: Dramatiq Broker
        :type broker: dramatiq.Broker
        :param worker: Dramatiq Worker
        :type worker: dramatiq.Worker
        """
        # 停止心跳线程
        self._running = False
        if self._heartbeat_thread:
            self._heartbeat_thread.join(timeout=5)

        # 注销 Worker
        self._unregister_worker()

        logger.info(f"Worker {self.worker_id[:8]} unregistered")

    def _register_worker(self) -> None:
        """
        注册 Worker 到 Redis
        """
        if not self.redis_client:
            logger.warning("Redis client not available, skipping worker registration")
            return

        now = int(time.time() * 1000)

        try:
            # 1. 添加到活跃 Worker 集合 (ZSET)
            self.redis_client.zadd("dramatiq:workers:active", {self.worker_id: now})

            # 2. 注册 Actor 列表 (SET)
            actor_key = f"dramatiq:worker:{self.worker_id}:actors"
            if self.actors:
                self.redis_client.sadd(actor_key, *self.actors)
                self.redis_client.expire(actor_key, self.worker_ttl)

            # 3. 注册 Worker 详细信息 (HASH)
            info_key = f"dramatiq:worker:{self.worker_id}:info"
            self.redis_client.hset(
                info_key,
                mapping={
                    "id": self.worker_id,
                    "hostname": self.hostname,
                    "pid": str(self.pid),
                    "started_at": str(self.started_at),
                    "last_heartbeat": str(now),
                    "actors": json.dumps(self.actors),
                },
            )
            self.redis_client.expire(info_key, self.worker_ttl)

            # 4. 注册 Actor -> Queue 映射 (HASH) - 用于动态队列发现
            actor_queues = {}
            for actor_name in self.actors:
                try:
                    actor = dramatiq.get_broker().get_actor(actor_name)
                    actor_queues[actor_name] = actor.queue_name
                except Exception as e:
                    logger.error(f"Failed to register actor: {e}")
                    pass

            if actor_queues:
                queue_map_key = "dramatiq:actor_queues"
                self.redis_client.hset(queue_map_key, mapping=actor_queues)

        except Exception as e:
            logger.error(f"Failed to register worker: {e}")

    def _unregister_worker(self) -> None:
        """
        从 Redis 注销 Worker
        """
        if not self.redis_client:
            return

        try:
            # 1. 从活跃集合移除
            self.redis_client.zrem("dramatiq:workers:active", self.worker_id)

            # 2. 删除 Actor 列表
            self.redis_client.delete(f"dramatiq:worker:{self.worker_id}:actors")

            # 3. 删除详细信息
            self.redis_client.delete(f"dramatiq:worker:{self.worker_id}:info")

        except Exception as e:
            logger.error(f"Failed to unregister worker: {e}")

    def _send_heartbeat(self) -> None:
        """
        发送心跳
        """
        if not self.redis_client:
            return

        now = int(time.time() * 1000)

        try:
            # 1. 更新活跃时间戳
            self.redis_client.zadd("dramatiq:workers:active", {self.worker_id: now})

            # 2. 更新心跳时间
            info_key = f"dramatiq:worker:{self.worker_id}:info"
            self.redis_client.hset(info_key, "last_heartbeat", str(now))
            self.redis_client.expire(info_key, self.worker_ttl)

            # 3. 刷新 Actor 列表过期时间
            actor_key = f"dramatiq:worker:{self.worker_id}:actors"
            self.redis_client.expire(actor_key, self.worker_ttl)

        except Exception as e:
            logger.error(f"Failed to send heartbeat: {e}")

    def _heartbeat_loop(self) -> None:
        """
        心跳循环（后台线程）
        """
        while self._running:
            time.sleep(self.heartbeat_interval)
            if self._running:
                self._send_heartbeat()
