import os
import time
import asyncio
import logging
import json
from typing import Dict, List, Optional, Set, Callable, Any

from .heartbeat import HeartbeatManager

logger = logging.getLogger(__name__)


# ============================================================================
# Worker 状态管理
# ============================================================================

class WorkerManager:
    """Worker管理器 - Worker状态的唯一管理入口

    ⚠️ 重要：所有Worker状态的修改都必须通过这个类进行，不要直接操作Redis！

    职责:
    1. Worker ID 生成和复用
    2. Worker 注册表管理
    3. Worker 状态字段的读写
    4. Redis Pub/Sub 状态变更通知
    5. Worker 查询和统计
    """

    def __init__(self, redis_client, async_redis_client, redis_prefix: str = "jettask", event_pool=None):
        """初始化Worker状态管理器

        Args:
            redis_client: 同步Redis客户端
            async_redis_client: 异步Redis客户端
            redis_prefix: Redis key前缀
            event_pool: EventPool实例（可选），用于事件驱动的消息恢复
        """
        self.sync_redis = redis_client
        self.redis = async_redis_client
        self.async_redis = async_redis_client
        self.redis_prefix = redis_prefix
        self.active_workers_key = f"{redis_prefix}:ACTIVE_WORKERS"
        self.workers_registry_key = f"{redis_prefix}:REGISTRY:WORKERS"
        self.event_pool = event_pool
        self.worker_prefix = 'WORKER'

        # Pub/Sub通道名称
        self.worker_state_channel = f"{redis_prefix}:WORKER_STATE_CHANGE"

        # 监听器订阅
        self._pubsub = None
        self._listener_task: Optional[asyncio.Task] = None
        self._running = False
        self._callbacks: Set[Callable] = set()

        # Pub/Sub 配置
        self._health_check_interval = 60
        self._health_check_task: Optional[asyncio.Task] = None

        # Worker扫描器相关（从 WorkerScanner 合并）
        self._initialized = False
        self._last_full_sync = 0
        self._full_sync_interval = 60
        self._scan_counter = 0
        self._partial_check_interval = 10

    def _get_worker_key(self, worker_id: str) -> str:
        """获取worker的Redis key"""
        return f"{self.redis_prefix}:WORKER:{worker_id}"

    async def initialize_worker(self, worker_id: str, worker_info: Dict[str, Any]):
        """初始化worker（首次创建）"""
        worker_key = self._get_worker_key(worker_id)
        current_time = time.time()

        worker_info.setdefault('is_alive', 'true')
        worker_info.setdefault('messages_transferred', 'false')
        worker_info.setdefault('created_at', str(current_time))
        worker_info.setdefault('last_heartbeat', str(current_time))

        pipeline = self.redis.pipeline()
        pipeline.hset(worker_key, mapping=worker_info)
        pipeline.zadd(self.active_workers_key, {worker_id: current_time})
        await pipeline.execute()

        logger.debug(f"Initialized worker {worker_id}")

    async def set_worker_online(self, worker_id: str, worker_data: dict = None):
        """设置worker为在线状态"""
        worker_key = self._get_worker_key(worker_id)
        old_alive = await self.redis.hget(worker_key, 'is_alive')
        old_alive = old_alive.decode('utf-8') if isinstance(old_alive, bytes) else old_alive

        current_time = time.time()
        pipeline = self.redis.pipeline()
        pipeline.hset(worker_key, 'is_alive', 'true')
        pipeline.hset(worker_key, 'last_heartbeat', str(current_time))

        # 当 worker 从离线变为在线时，重置 messages_transferred
        # 这表示是一个新的 worker 实例，还没有进行消息转移
        if old_alive != 'true':
            pipeline.hset(worker_key, 'messages_transferred', 'false')

        if worker_data:
            pipeline.hset(worker_key, mapping=worker_data)

        pipeline.zadd(self.active_workers_key, {worker_id: current_time})
        await pipeline.execute()

        if old_alive != 'true':
            await self._publish_state_change(worker_id, 'online')
            logger.debug(f"Worker {worker_id} is now ONLINE")

    async def set_worker_offline(self, worker_id: str, reason: str = "unknown"):
        """设置worker为离线状态"""
        worker_key = self._get_worker_key(worker_id)
        old_alive = await self.redis.hget(worker_key, 'is_alive')
        old_alive = old_alive.decode('utf-8') if isinstance(old_alive, bytes) else old_alive

        current_time = time.time()
        pipeline = self.redis.pipeline()
        pipeline.hset(worker_key, 'messages_transferred', 'false')  # 重置消息转移标记，允许其他worker接管消息
        pipeline.hset(worker_key, 'is_alive', 'false')
        pipeline.hset(worker_key, 'offline_reason', reason)
        pipeline.hset(worker_key, 'offline_time', str(current_time))
        pipeline.zrem(self.active_workers_key, worker_id)
        await pipeline.execute()

        if old_alive == 'true':
            await self._publish_state_change(worker_id, 'offline', reason)
            logger.debug(f"Worker {worker_id} is now OFFLINE (reason: {reason})")

    async def update_worker_heartbeat(self, worker_id: str, heartbeat_data: dict = None):
        """更新worker心跳（确保在线状态）"""
        worker_key = self._get_worker_key(worker_id)
        current_time = time.time()

        pipeline = self.redis.pipeline()
        pipeline.hset(worker_key, 'is_alive', 'true')
        pipeline.hset(worker_key, 'last_heartbeat', str(current_time))

        if heartbeat_data:
            pipeline.hset(worker_key, mapping=heartbeat_data)

        pipeline.zadd(self.active_workers_key, {worker_id: current_time})
        await pipeline.execute()

    async def update_worker_field(self, worker_id: str, field: str, value: str):
        """更新worker的单个字段"""
        worker_key = self._get_worker_key(worker_id)
        await self.redis.hset(worker_key, field, value)

    async def update_worker_fields(self, worker_id: str, fields: Dict[str, Any]):
        """批量更新worker的多个字段"""
        worker_key = self._get_worker_key(worker_id)
        await self.redis.hset(worker_key, mapping=fields)

    async def record_group_info(
        self,
        worker_id: str,
        queue: str,
        task_name: str,
        group_name: str,
        consumer_name: str,
        redis_prefix: str
    ):
        """记录 task 的 group 信息到 worker hash 表

        Args:
            worker_id: Worker ID
            queue: 队列名（不带前缀）
            task_name: 任务名
            group_name: Consumer group 名称
            consumer_name: Consumer 名称
            redis_prefix: Redis 前缀
        """
        try:
            worker_key = self._get_worker_key(worker_id)

            # 构建 group 信息
            group_info = {
                'queue': queue,
                'task_name': task_name,
                'group_name': group_name,
                'consumer_name': consumer_name,
                'stream_key': f"{redis_prefix}:QUEUE:{queue}"
            }

            # 将 group 信息存储到 worker 的 hash 中
            field_name = f"group_info:{group_name}"
            await self.redis.hset(
                worker_key,
                field_name,
                json.dumps(group_info)
            )
            logger.debug(f"Recorded group info for worker {worker_id}, task {task_name}: {group_info}")

        except Exception as e:
            logger.error(f"Error recording group info for worker {worker_id}: {e}", exc_info=True)

    def update_worker_field_sync(self, worker_id: str, field: str, value: str):
        """更新worker的单个字段（同步版本，供心跳线程使用）"""
        worker_key = self._get_worker_key(worker_id)
        self.sync_redis.hset(worker_key, field, value)
        logger.debug(f"Updated worker {worker_id} field {field}={value}")

    def update_worker_fields_sync(self, worker_id: str, fields: Dict[str, Any]):
        """批量更新worker的多个字段（同步版本，供心跳线程使用）"""
        worker_key = self._get_worker_key(worker_id)
        self.sync_redis.hset(worker_key, mapping=fields)
        logger.debug(f"Updated worker {worker_id} fields: {list(fields.keys())}")

    async def increment_queue_stats(self, worker_id: str, queue: str,
                                   running_tasks_delta: int = None,
                                   success_count_increment: int = None,
                                   failed_count_increment: int = None,
                                   total_count_increment: int = None,
                                   processing_time_increment: float = None,
                                   latency_time_increment: float = None):
        """增量更新worker在特定队列上的累积统计信息"""
        worker_key = self._get_worker_key(worker_id)
        pipeline = self.redis.pipeline()

        if running_tasks_delta is not None and running_tasks_delta != 0:
            pipeline.hincrby(worker_key, f'{queue}:running_tasks', running_tasks_delta)

        if success_count_increment is not None:
            pipeline.hincrby(worker_key, f'{queue}:success_count', success_count_increment)

        if failed_count_increment is not None:
            pipeline.hincrby(worker_key, f'{queue}:failed_count', failed_count_increment)

        if total_count_increment is not None:
            pipeline.hincrby(worker_key, f'{queue}:total_count', total_count_increment)

        if processing_time_increment is not None:
            pipeline.hincrbyfloat(worker_key, f'{queue}:total_processing_time', processing_time_increment)

        if latency_time_increment is not None:
            pipeline.hincrbyfloat(worker_key, f'{queue}:total_latency_time', latency_time_increment)

        await pipeline.execute()

    async def get_queue_total_stats(self, worker_id: str, queue: str) -> dict:
        """获取队列的累积统计数据"""
        worker_key = self._get_worker_key(worker_id)
        fields = [
            f'{queue}:total_count',
            f'{queue}:total_processing_time',
            f'{queue}:total_latency_time'
        ]
        values = await self.redis.hmget(worker_key, fields)

        return {
            'total_count': int(values[0]) if values[0] else 0,
            'total_processing_time': float(values[1]) if values[1] else 0.0,
            'total_latency_time': float(values[2]) if values[2] else 0.0
        }

    async def update_queue_stats(self, worker_id: str, queue: str,
                                 running_tasks: int = None,
                                 avg_processing_time: float = None,
                                 avg_latency_time: float = None):
        """更新worker在特定队列上的统计信息"""
        worker_key = self._get_worker_key(worker_id)
        pipeline = self.redis.pipeline()

        if running_tasks is not None:
            pipeline.hset(worker_key, f'{queue}:running_tasks', str(running_tasks))

        if avg_processing_time is not None:
            pipeline.hset(worker_key, f'{queue}:avg_processing_time', f'{avg_processing_time:.3f}')

        if avg_latency_time is not None:
            pipeline.hset(worker_key, f'{queue}:avg_latency_time', f'{avg_latency_time:.3f}')

        await pipeline.execute()

    async def mark_messages_transferred(self, worker_id: str, transferred: bool = True):
        """标记worker的消息是否已转移"""
        worker_key = self._get_worker_key(worker_id)
        await self.redis.hset(worker_key, 'messages_transferred', 'true' if transferred else 'false')

    async def get_worker_info(self, worker_id: str) -> Optional[Dict[str, str]]:
        """获取worker的完整信息"""
        worker_key = self._get_worker_key(worker_id)
        data = await self.redis.hgetall(worker_key)

        if not data:
            return None

        result = {}
        for k, v in data.items():
            key = k.decode('utf-8') if isinstance(k, bytes) else k
            value = v.decode('utf-8') if isinstance(v, bytes) else v
            result[key] = value

        return result

    async def get_worker_field(self, worker_id: str, field: str) -> Optional[str]:
        """获取worker的单个字段值"""
        worker_key = self._get_worker_key(worker_id)
        value = await self.redis.hget(worker_key, field)

        if value is None:
            return None

        return value.decode('utf-8') if isinstance(value, bytes) else value

    async def is_worker_alive(self, worker_id: str) -> bool:
        """检查worker是否在线"""
        is_alive = await self.get_worker_field(worker_id, 'is_alive')
        return is_alive == 'true'

    # ========== Worker ID 生成和复用 ==========

    def generate_worker_id(self, prefix: str) -> str:
        """
        生成新的 Worker ID

        格式: {prefix}-{uuid}-{pid}
        例如: YYDG-a1b2c3d4-12345

        Args:
            prefix: Worker ID 前缀（通常是主机名）

        Returns:
            生成的 Worker ID
        """
        import uuid
        return f"{prefix}-{uuid.uuid4().hex[:8]}-{os.getpid()}"

    async def find_reusable_worker_id(self, prefix: str) -> Optional[str]:
        """
        查找可复用的离线 Worker ID

        Args:
            prefix: Worker ID 前缀

        Returns:
            可复用的 Worker ID，如果没有则返回 None
        """
        try:
            offline_workers = await self.get_offline_workers()

            for worker_id in offline_workers:
                if isinstance(worker_id, bytes):
                    worker_id = worker_id.decode('utf-8')
                if worker_id.startswith(prefix):
                    logger.debug(f"Found reusable worker ID: {worker_id}")
                    return worker_id
        except Exception as e:
            logger.warning(f"Error finding reusable worker ID: {e}")

        return None

    # ========== Worker 注册表管理 ==========

    async def register_worker(self, worker_id: str):
        """注册 Worker 到全局注册表"""
        await self.async_redis.sadd(self.workers_registry_key, worker_id)
        logger.debug(f"Registered worker: {worker_id}")

    async def unregister_worker(self, worker_id: str):
        """从全局注册表注销 Worker"""
        await self.async_redis.srem(self.workers_registry_key, worker_id)
        logger.debug(f"Unregistered worker: {worker_id}")

    async def get_all_workers(self) -> Set[str]:
        """获取所有已注册的 Worker ID"""
        return await self.async_redis.smembers(self.workers_registry_key)

    def get_all_workers_sync(self) -> Set[str]:
        """同步方式获取所有已注册的 Worker ID"""
        return self.sync_redis.smembers(self.workers_registry_key)

    async def get_worker_count(self) -> int:
        """获取已注册的 Worker 总数"""
        return await self.async_redis.scard(self.workers_registry_key)

    async def get_offline_workers(self) -> Set[str]:
        """获取所有离线的 Worker ID

        离线 Worker 是指已注册但 is_alive=false 的 Worker
        """
        all_workers = await self.get_all_workers()
        offline_workers = set()

        for worker_id in all_workers:
            if isinstance(worker_id, bytes):
                worker_id = worker_id.decode('utf-8')

            worker_key = f"{self.redis_prefix}:WORKER:{worker_id}"
            is_alive = await self.async_redis.hget(worker_key, 'is_alive')

            if is_alive:
                is_alive = is_alive.decode('utf-8') if isinstance(is_alive, bytes) else is_alive
                if is_alive != 'true':
                    offline_workers.add(worker_id)
            else:
                # Worker key 不存在或没有 is_alive 字段，认为离线
                offline_workers.add(worker_id)

        return offline_workers

    async def get_workers_for_task(self, task_name: str, only_alive: bool = True) -> Set[str]:
        """获取执行特定任务的 Worker 列表

        通过检查 WORKER:* hash 中的 group_info 字段来判断哪些 Worker 在处理该任务

        Args:
            task_name: 任务名称
            only_alive: 是否只返回在线的 Worker（默认 True）

        Returns:
            处理该任务的 Worker ID 集合
        """
        all_worker_ids = await self.get_all_workers()
        matched_workers = set()
        group_info_prefix = f"group_info:{self.redis_prefix}:QUEUE:"

        for worker_id in all_worker_ids:
            if isinstance(worker_id, bytes):
                worker_id = worker_id.decode('utf-8')

            worker_key = f"{self.redis_prefix}:WORKER:{worker_id}"
            worker_info = await self.async_redis.hgetall(worker_key)

            if not worker_info:
                continue

            # 解码 bytes keys
            decoded_info = {}
            for k, v in worker_info.items():
                key = k.decode('utf-8') if isinstance(k, bytes) else k
                val = v.decode('utf-8') if isinstance(v, bytes) else v
                decoded_info[key] = val

            # 检查 is_alive 状态
            if only_alive:
                is_alive = decoded_info.get('is_alive', 'false')
                if is_alive != 'true':
                    continue

            # 检查是否包含该任务的 group_info
            # 格式: group_info:test5:QUEUE:robust_bench2:benchmark_task
            for key in decoded_info.keys():
                if key.startswith(group_info_prefix):
                    parts = key.split(':')
                    if len(parts) >= 5:
                        worker_task_name = parts[-1]  # 最后一部分是 task_name
                        if worker_task_name == task_name:
                            matched_workers.add(worker_id)
                            break

        return matched_workers

    async def get_active_worker_count_for_task(self, task_name: str) -> int:
        """获取执行特定任务的在线 Worker 数量

        Args:
            task_name: 任务名称

        Returns:
            在线 Worker 数量
        """
        workers = await self.get_workers_for_task(task_name, only_alive=True)
        return len(workers)

    async def find_offline_workers_for_task(
        self,
        task_name: str,
        worker_prefix: str = 'WORKER'
    ) -> List[tuple]:
        """查找指定任务的离线 Worker

        查找条件：
        1. Worker 已离线（is_alive=false）
        2. 消息未转移（messages_transferred=false）
        3. Worker 的 group_info 中有该任务

        Args:
            task_name: 任务名称
            worker_prefix: Worker 键前缀（默认 'WORKER'）

        Returns:
            离线 Worker 列表，每项为 (worker_key, worker_data) 元组
        """
        offline_workers = []

        try:
            # 获取所有 worker ID
            worker_ids = await self.get_all_workers()
            logger.debug(f"[Recovery] Found {len(worker_ids)} workers in registry for task {task_name}")

            for worker_id in worker_ids:
                # 解码 worker_id（可能是 bytes）
                if isinstance(worker_id, bytes):
                    worker_id = worker_id.decode('utf-8')

                # 构建 worker key
                worker_key = f"{self.redis_prefix}:{worker_prefix}:{worker_id}"

                try:
                    # 直接使用自身的方法读取 worker 信息
                    decoded_worker_data = await self.get_worker_info(worker_id)

                    if not decoded_worker_data:
                        continue

                    # 检查 worker 是否离线且消息未转移
                    is_alive = decoded_worker_data.get('is_alive', 'false') == 'true'
                    messages_transferred = decoded_worker_data.get('messages_transferred', 'false') == 'true'

                    # 找到离线且消息未转移的 worker
                    if not is_alive and not messages_transferred:
                        # 检查 worker 的 group_info 中是否有该任务
                        has_task = False
                        for key, value in decoded_worker_data.items():
                            if key.startswith('group_info:'):
                                try:
                                    import json
                                    group_info = json.loads(value)
                                    if group_info.get('task_name') == task_name:
                                        has_task = True
                                        break
                                except Exception as e:
                                    logger.error(f"Error parsing group_info: {e}")

                        if has_task:
                            logger.info(
                                f"Found offline worker needing recovery: {worker_id}, "
                                f"task={task_name}, is_alive={is_alive}, "
                                f"messages_transferred={messages_transferred}"
                            )
                            offline_workers.append((worker_key, decoded_worker_data))
                        else:
                            logger.debug(
                                f"Worker {worker_id} is offline but not responsible for task {task_name}"
                            )

                except Exception as e:
                    logger.error(f"Error processing worker key {worker_key}: {e}")
                    continue

        except Exception as e:
            logger.error(f"Error finding offline workers: {e}")

        return offline_workers

    async def get_all_workers_info(self, only_alive: bool = True) -> Dict[str, Dict[str, str]]:
        """获取所有worker的信息"""
        pattern = f"{self.redis_prefix}:WORKER:*"
        result = {}

        cursor = 0
        while True:
            cursor, keys = await self.redis.scan(cursor, match=pattern, count=100)

            for key in keys:
                if isinstance(key, bytes):
                    key = key.decode('utf-8')

                parts = key.split(":")
                if len(parts) >= 3:
                    worker_id = parts[2]
                    worker_info = await self.get_worker_info(worker_id)
                    if worker_info:
                        if only_alive and worker_info.get('is_alive') != 'true':
                            continue
                        result[worker_id] = worker_info

            if cursor == 0:
                break

        return result

    async def delete_worker(self, worker_id: str):
        """删除worker的所有数据"""
        worker_key = self._get_worker_key(worker_id)
        pipeline = self.redis.pipeline()
        pipeline.delete(worker_key)
        pipeline.zrem(self.active_workers_key, worker_id)
        await pipeline.execute()
        logger.debug(f"Deleted worker {worker_id}")

    async def _publish_state_change(self, worker_id: str, state: str, reason: str = None):
        """发布状态变更信号"""
        message = {
            'worker_id': worker_id,
            'state': state,
            'timestamp': asyncio.get_event_loop().time()
        }

        if reason:
            message['reason'] = reason

        await self.redis.publish(
            self.worker_state_channel,
            json.dumps(message)
        )

        logger.debug(f"Published state change: {message}")

    async def start_listener(self):
        """启动状态变更监听器"""
        if self._running:
            logger.warning("Worker state listener already running")
            return

        self._running = True
        self._pubsub = await self._create_and_subscribe_pubsub()
        self._listener_task = asyncio.create_task(self._listen_loop())
        self._health_check_task = asyncio.create_task(self._health_check_loop())

        logger.debug(f"Started worker state listener on channel: {self.worker_state_channel}")

    async def stop_listener(self):
        """停止状态变更监听器"""
        if not self._running:
            return

        self._running = False

        if self._listener_task:
            self._listener_task.cancel()
            try:
                await self._listener_task
            except asyncio.CancelledError:
                pass

        if self._health_check_task:
            self._health_check_task.cancel()
            try:
                await self._health_check_task
            except asyncio.CancelledError:
                pass

        if self._pubsub:
            await self._pubsub.unsubscribe(self.worker_state_channel)
            await self._pubsub.close()

        logger.debug("Stopped worker state listener")

    async def _create_and_subscribe_pubsub(self):
        """创建 PubSub 连接并订阅频道"""
        if self._pubsub:
            try:
                await self._pubsub.close()
            except:
                pass

        pubsub = self.redis.pubsub()
        await pubsub.subscribe(self.worker_state_channel)

        # 标记PubSub连接，防止被空闲连接清理
        if hasattr(pubsub, 'connection') and pubsub.connection:
            pubsub.connection._is_pubsub_connection = True
            socket_timeout = pubsub.connection.socket_timeout if hasattr(pubsub.connection, 'socket_timeout') else 'N/A'
            logger.info(f"Marked PubSub connection {id(pubsub.connection)} to prevent cleanup, socket_timeout={socket_timeout}")

        logger.debug(f"Created and subscribed to Redis Pub/Sub channel: {self.worker_state_channel}")
        return pubsub

    async def _health_check_loop(self):
        """定期检查 Pub/Sub 连接健康状态"""
        while self._running:
            try:
                await asyncio.sleep(self._health_check_interval)

                if not self._running:
                    break

                if self._pubsub and self._pubsub.connection:
                    try:
                        await asyncio.wait_for(self._pubsub.ping(), timeout=5.0)
                        logger.debug("Pub/Sub health check: OK")
                    except Exception as e:
                        logger.warning(f"Pub/Sub health check failed: {e}")
                else:
                    logger.warning("Pub/Sub connection is None")

            except asyncio.CancelledError:
                logger.debug("Health check loop cancelled")
                break
            except Exception as e:
                logger.error(f"Error in health check loop: {e}")

    async def _listen_loop(self):
        """监听循环（支持自动重连）"""
        retry_delay = 1
        max_retry_delay = 30

        while self._running:
            try:
                async for message in self._pubsub.listen():
                    if message['type'] == 'message':
                        try:
                            data = json.loads(message['data'])

                            if data.get('state') == 'offline' and self.event_pool:
                                worker_id = data.get('worker_id')
                                if worker_id:
                                    logger.info(f"[StateManager] Worker {worker_id} offline event received")
                                    asyncio.create_task(
                                        self.event_pool.handle_worker_offline_event(worker_id)
                                    )

                            for callback in self._callbacks:
                                try:
                                    if asyncio.iscoroutinefunction(callback):
                                        await callback(data)
                                    else:
                                        callback(data)
                                except Exception as e:
                                    logger.error(f"Error in state change callback: {e}")

                        except Exception as e:
                            logger.error(f"Error processing state change message: {e}")

                retry_delay = 1

            except asyncio.CancelledError:
                logger.debug("Listen loop cancelled")
                break
            except Exception as e:
                logger.error(f"Error in listen loop: {e}")

                if not self._running:
                    break

                logger.warning(f"Attempting to reconnect to Redis Pub/Sub in {retry_delay} seconds...")
                await asyncio.sleep(retry_delay)

                try:
                    self._pubsub = await self._create_and_subscribe_pubsub()
                    logger.info(f"Successfully reconnected to Redis Pub/Sub")
                    retry_delay = 1
                except Exception as reconnect_error:
                    logger.error(f"Failed to reconnect to Redis Pub/Sub: {reconnect_error}")
                    retry_delay = min(retry_delay * 2, max_retry_delay)

        logger.debug("Listen loop exited")

    def register_callback(self, callback: Callable):
        """注册状态变更回调"""
        self._callbacks.add(callback)
        logger.debug(f"Registered state change callback: {callback.__name__}")

    def unregister_callback(self, callback: Callable):
        """注销状态变更回调"""
        self._callbacks.discard(callback)
        logger.debug(f"Unregistered state change callback: {callback.__name__}")

    # ========================================================================
    # Worker 扫描和超时检测（从 WorkerScanner 合并）
    # ========================================================================

    async def scan_timeout_workers(self) -> List[Dict]:
        """快速扫描超时的 worker - O(log N) 复杂度

        每个 worker 使用自己存储的 heartbeat_timeout 值进行超时判断。
        这个值由 HeartbeatManager 在 worker 初始化时设置。

        Returns:
            超时的 worker 列表
        """
        self._scan_counter += 1
        if self._scan_counter >= self._partial_check_interval:
            self._scan_counter = 0
            asyncio.create_task(self._partial_check())

        current_time = time.time()

        potential_timeout_worker_ids = await self.async_redis.zrangebyscore(
            self.active_workers_key,
            min=0,
            max=current_time - 1
        )

        if not potential_timeout_worker_ids:
            return []

        # 直接使用自身的方法获取 worker 信息
        all_workers_info = await self.get_all_workers_info(only_alive=False)
        workers_data = [all_workers_info.get(wid) for wid in potential_timeout_worker_ids]

        result = []
        cleanup_pipeline = self.async_redis.pipeline()
        need_cleanup = False

        for worker_id, worker_data in zip(potential_timeout_worker_ids, workers_data):
            if not worker_data:
                cleanup_pipeline.zrem(self.active_workers_key, worker_id)
                cleanup_pipeline.srem(self.workers_registry_key, worker_id)
                need_cleanup = True
                continue

            # 使用 worker 自己存储的 heartbeat_timeout
            # 这个值在 HeartbeatManager 初始化 worker 时设置
            worker_heartbeat_timeout = float(worker_data.get('heartbeat_timeout', 15.0))
            last_heartbeat = float(worker_data.get('last_heartbeat', 0))
            worker_cutoff_time = current_time - worker_heartbeat_timeout

            if last_heartbeat >= worker_cutoff_time:
                cleanup_pipeline.zadd(self.active_workers_key, {worker_id: last_heartbeat})
                need_cleanup = True
                continue

            is_alive = worker_data.get('is_alive', 'true') == 'true'
            if not is_alive:
                cleanup_pipeline.zrem(self.active_workers_key, worker_id)
                need_cleanup = True
                continue

            logger.debug(f"Worker {worker_id} timeout: last_heartbeat={last_heartbeat}, timeout={worker_heartbeat_timeout}s")
            worker_key = f"{self.redis_prefix}:{self.worker_prefix}:{worker_id}"
            result.append({
                'worker_key': worker_key,
                'worker_data': worker_data,
                'worker_id': worker_id
            })

        if need_cleanup:
            await cleanup_pipeline.execute()

        if result:
            logger.info(f"Found {len(result)} timeout workers")

        return result

    async def update_heartbeat(self, worker_id: str, heartbeat_time: Optional[float] = None):
        """原子性更新心跳"""
        if heartbeat_time is None:
            heartbeat_time = time.time()

        pipeline = self.async_redis.pipeline()
        worker_key = f"{self.redis_prefix}:{self.worker_prefix}:{worker_id}"

        pipeline.hset(worker_key, 'last_heartbeat', str(heartbeat_time))
        pipeline.zadd(self.active_workers_key, {worker_id: heartbeat_time})

        await pipeline.execute()

    async def add_worker(self, worker_id: str, worker_data: Dict):
        """添加新 worker"""
        heartbeat_time = float(worker_data.get('last_heartbeat', time.time()))

        pipeline = self.async_redis.pipeline()
        worker_key = f"{self.redis_prefix}:{self.worker_prefix}:{worker_id}"

        pipeline.hset(worker_key, mapping=worker_data)
        pipeline.zadd(self.active_workers_key, {worker_id: heartbeat_time})

        await pipeline.execute()
        logger.debug(f"Added worker {worker_id} to system")

    async def remove_worker(self, worker_id: str):
        """移除 worker（标记为离线并从活跃集合中移除）"""
        await self.set_worker_offline(worker_id, reason="heartbeat_timeout")
        await self.async_redis.zrem(self.active_workers_key, worker_id)

    async def cleanup_stale_workers(self, max_age_seconds: float = 3600):
        """清理过期的 worker 记录"""
        current_time = time.time()
        cutoff_time = current_time - max_age_seconds

        stale_worker_ids = await self.async_redis.zrangebyscore(
            self.active_workers_key,
            min=0,
            max=cutoff_time
        )

        if not stale_worker_ids:
            return 0

        pipeline = self.async_redis.pipeline()

        for worker_id in stale_worker_ids:
            worker_key = f"{self.redis_prefix}:{self.worker_prefix}:{worker_id}"
            pipeline.delete(worker_key)

        pipeline.zrem(self.active_workers_key, *stale_worker_ids)

        await pipeline.execute()

        logger.info(f"Cleaned up {len(stale_worker_ids)} stale worker records")
        return len(stale_worker_ids)

    async def _partial_check(self):
        """部分一致性检查"""
        try:
            sample_size = min(10, await self.async_redis.zcard(self.active_workers_key))
            if sample_size == 0:
                return

            random_workers = await self.async_redis.zrandmember(
                self.active_workers_key, sample_size, withscores=True
            )

            for worker_id, zset_score in random_workers:
                worker_key = f"{self.redis_prefix}:{self.worker_prefix}:{worker_id}"
                hash_heartbeat = await self.async_redis.hget(worker_key, 'last_heartbeat')

                if not hash_heartbeat:
                    await self.async_redis.zrem(self.active_workers_key, worker_id)
                    logger.debug(f"Partial check: removed {worker_id}")
                else:
                    hash_time = float(hash_heartbeat)
                    if abs(hash_time - zset_score) > 1.0:
                        await self.async_redis.zadd(self.active_workers_key, {worker_id: hash_time})
                        logger.debug(f"Partial check: synced {worker_id}")

        except Exception as e:
            logger.debug(f"Partial check error: {e}")

    async def get_active_count(self) -> int:
        """获取活跃 worker 数量 - O(1)"""
        return await self.async_redis.zcard(self.active_workers_key)


__all__ = [
    'WorkerManager',
    'HeartbeatManager',
]
