import os
import time
import asyncio
import logging
import json
from typing import Dict, List, Optional, Set, Callable, Any
import msgpack
from redis.asyncio.lock import Lock as AsyncLock
from jettask.db.pubsub import PubSubManager, PubSubSubscription

logger = logging.getLogger(__name__)



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,
                 queue_formatter=None, queue_registry=None, app=None, tasks: dict = None,
                 task_event_queues: dict = None, worker_id: str = None):
        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'

        self.worker_state_channel = f"{redis_prefix}:WORKER_STATE_CHANGE"

        self._pubsub_subscription: Optional[PubSubSubscription] = None
        self._running = False
        self._callbacks: Set[Callable] = set()

        self._initialized = False
        self._last_full_sync = 0
        self._full_sync_interval = 60
        self._scan_counter = 0
        self._partial_check_interval = 10

        self._stop_recovery = False
        self._recovery_task = None  
        self.queue_formatter = queue_formatter or (lambda q: f"{self.redis_prefix}:QUEUE:{q}")
        self.worker_state_manager = app.worker_state_manager if (app and hasattr(app, 'worker_state_manager')) else None
        self.queue_registry = queue_registry
        self.app = app
        self.tasks = tasks or {}
        self.task_event_queues = task_event_queues or {}
        self.worker_id = worker_id

        self._stop_cleanup = False
        self._cleanup_task = None  

    def _get_worker_key(self, worker_id: str) -> str:
        return f"{self.redis_prefix}:WORKER:{worker_id}"

    async def initialize_worker(self, worker_id: str, worker_info: Dict[str, Any]):
        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_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))

        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_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')  
        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_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_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_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
    ):
        try:
            worker_key = self._get_worker_key(worker_id)

            group_info = {
                'queue': queue,
                'task_name': task_name,
                'group_name': group_name,
                'consumer_name': consumer_name,
                'stream_key': f"{redis_prefix}:QUEUE:{queue}"
            }

            field_name = f"group_info:{queue}"
            await self.redis.hset(
                worker_key,
                field_name,
                json.dumps(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_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_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_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_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_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_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_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:
        is_alive = await self.get_worker_field(worker_id, 'is_alive')
        return is_alive == 'true'


    def generate_worker_id(self, prefix: str) -> str:
        import uuid
        return f"{prefix}-{uuid.uuid4().hex[:8]}-{os.getpid()}"

    async def find_reusable_worker_id(self, prefix: str) -> Optional[str]:
        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.debug(f"Error finding reusable worker ID: {e}")

        return None


    async def register_worker(self, worker_id: str):
        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):
        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]:
        return await self.async_redis.smembers(self.workers_registry_key)

    def get_all_workers_sync(self) -> Set[str]:
        return self.sync_redis.smembers(self.workers_registry_key)

    async def get_worker_count(self) -> int:
        return await self.async_redis.scard(self.workers_registry_key)

    async def get_offline_workers(self) -> Set[str]:
        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:
                offline_workers.add(worker_id)

        return offline_workers

    async def get_workers_for_task(self, task_name: str, only_alive: bool = True) -> Set[str]:
        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

            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

            if only_alive:
                is_alive = decoded_info.get('is_alive', 'false')
                if is_alive != 'true':
                    continue

            for key in decoded_info.keys():
                if key.startswith(group_info_prefix):
                    parts = key.split(':')
                    if len(parts) >= 5:
                        worker_task_name = parts[-1]  
                        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:
        workers = await self.get_workers_for_task(task_name, only_alive=True)
        return len(workers)

    async def find_all_offline_workers(
        self,
        worker_prefix: str = 'WORKER'
    ):
        try:
            worker_ids = await self.get_all_workers()
            logger.debug(f"[Recovery] Scanning {len(worker_ids)} workers in registry")

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

                worker_key = f"{self.redis_prefix}:{worker_prefix}:{worker_id}"

                try:
                    decoded_worker_data = await self.get_worker_info(worker_id)

                    if not decoded_worker_data:
                        continue

                    is_alive = decoded_worker_data.get('is_alive', 'false') == 'true'
                    messages_transferred = decoded_worker_data.get('messages_transferred', 'false') == 'true'

                    if not is_alive and not messages_transferred:
                        logger.debug(
                            f"[Recovery] Found offline worker: {worker_id}, "
                            f"is_alive={is_alive}, "
                            f"messages_transferred={messages_transferred}"
                        )
                        yield (worker_key, decoded_worker_data)

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

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

    async def find_offline_workers_for_task(
        self,
        task_name: str,
        worker_prefix: str = 'WORKER'
    ) -> List[tuple]:
        all_offline_workers = await self.find_all_offline_workers(worker_prefix)
        logger.debug(f"[Recovery] Found {len(all_offline_workers)} offline workers, filtering for task {task_name}")

        task_offline_workers = []

        for worker_key, worker_data in all_offline_workers:
            has_task = False
            for key, value in worker_data.items():
                if key.startswith('group_info:'):
                    try:
                        group_info = json.loads(value)
                        if group_info.get('task_name') == task_name:
                            has_task = True
                            break
                    except Exception as e:
                        logger.error(f"[Recovery] Error parsing group_info: {e}")

            if has_task:
                worker_id = worker_key.split(':')[-1]
                logger.debug(
                    f"[Recovery] Found offline worker needing recovery: {worker_id}, task={task_name}"
                )
                task_offline_workers.append((worker_key, worker_data))
            else:
                worker_id = worker_key.split(':')[-1]
                logger.debug(
                    f"[Recovery] Worker {worker_id} is offline but not responsible for task {task_name}"
                )

        logger.debug(f"[Recovery] Found {len(task_offline_workers)} offline workers for task {task_name}")
        return task_offline_workers

    async def get_all_workers_info(self, only_alive: bool = True) -> Dict[str, Dict[str, str]]:
        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, unregister: bool = True):
        worker_key = self._get_worker_key(worker_id)
        pipeline = self.redis.pipeline()
        pipeline.delete(worker_key)  
        pipeline.zrem(self.active_workers_key, worker_id)  

        if unregister:
            pipeline.srem(self.workers_registry_key, worker_id)  

        await pipeline.execute()
        logger.debug(f"Deleted worker {worker_id}" + (" and unregistered" if unregister else ""))

    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.debug("Worker state listener already running")
            return

        self._running = True

        self._pubsub_subscription = await PubSubManager.subscribe(
            redis_client=self.redis,
            channel=self.worker_state_channel,
            on_message=self.handle_state_change,
            name="WorkerStateListener",
        )

        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._pubsub_subscription:
            await self._pubsub_subscription.stop()
            self._pubsub_subscription = None

        logger.debug("Stopped worker state listener")

    async def handle_state_change(self, data: dict):
        try:
            if data.get('state') == 'offline':
                worker_id = data.get('worker_id')
                if worker_id:
                    logger.debug(f"[StateManager] Worker {worker_id} offline event received, triggering recovery")
                    try:
                        await self._process_offline_workers_once()
                    except Exception as e:
                        logger.error(f"[StateManager] Error during offline worker recovery: {e}", exc_info=True)

            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}")

    _handle_state_change_message = handle_state_change

    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__}")


    async def scan_timeout_workers(self) -> List[Dict]:
        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 []

        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 = 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.debug(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):
        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):
        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):
        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.debug(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
            )

            if random_workers:
                worker_pairs = [(random_workers[i], random_workers[i+1]) for i in range(0, len(random_workers), 2)]
            else:
                worker_pairs = []

            for worker_id, zset_score in worker_pairs:
                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)
                    zset_score_float = float(zset_score)
                    if abs(hash_time - zset_score_float) > 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:
            import traceback
            traceback.print_exc()
            logger.error(f"Partial check error: {e}", exc_info=True)

    async def get_active_count(self) -> int:
        return await self.async_redis.zcard(self.active_workers_key)


    async def start_recovery(self, recovery_interval: float = 30.0):
        logger.debug(f"[Recovery] Starting offline worker recovery loop (interval={recovery_interval}s)")
        self._stop_recovery = False

        while not self._stop_recovery:
            try:
                await self._process_offline_workers_once()

                await asyncio.sleep(recovery_interval)

            except asyncio.CancelledError:
                logger.debug("[Recovery] Recovery loop cancelled")
                break
            except Exception as e:
                logger.error(f"[Recovery Loop] Unexpected error: {e}")
                await asyncio.sleep(recovery_interval)

        logger.debug("[Recovery] Offline worker recovery loop stopped")

    async def start(self, recovery_interval: float = 30.0):
        await self.start_recovery(recovery_interval)

    async def _process_offline_workers_once(self):
        

        async for worker_key, worker_data in self.find_all_offline_workers():
            
            if self._stop_recovery:
                break

            worker_id = worker_key.split(':')[-1]

            lock_key = f"{self.redis_prefix}:RECOVERY:WORKER_LOCK:{worker_id}"
            worker_lock = AsyncLock(
                self.async_redis,
                lock_key,
                timeout=300,  
                blocking=False  
            )

            if not await worker_lock.acquire():
                logger.debug(
                    f"[Recovery] Worker {worker_id} is being processed by another worker, skipping"
                )
                continue

            try:
                logger.debug(f"[Recovery] Processing worker={worker_id}")

                for key, value in worker_data.items():
                    if not key.startswith('group_info:'):
                        continue

                    try:
                        group_info = json.loads(value)
                        task_name = group_info.get('task_name')

                        already_transferred = group_info.get('messages_transferred', 'false')
                        if already_transferred == 'true':
                            logger.debug(
                                f"[Recovery] Group_info already transferred: worker={worker_id}, "
                                f"task={task_name}, skipping"
                            )
                            continue

                        if task_name not in self.tasks:
                            logger.debug(f"[Recovery] Task {task_name} not in managed tasks, skipping {self.tasks=}")
                            continue

                        task_event_queue = self.task_event_queues.get(task_name)
                        if not task_event_queue:
                            logger.debug(f"[Recovery] No event queue for task {task_name}, skipping")
                            continue

                        queue_name = group_info.get('queue')  
                        group_name = group_info.get('group_name')
                        offline_consumer_name = group_info.get('consumer_name')
                        stream_key_suffix = group_info.get('stream_key')  

                        if not all([queue_name, group_name, offline_consumer_name]):
                            logger.debug(f"[Recovery] Incomplete group_info: {group_info}")
                            continue

                        stream_key = stream_key_suffix

                        logger.debug(
                            f"[Recovery] Processing worker={worker_id}, task={task_name}, "
                            f"queue={queue_name}, group={group_name}"
                        )

                        success = await self._transfer_messages(
                            stream_key=stream_key,
                            group_name=group_name,
                            offline_consumer_name=offline_consumer_name,
                            task_name=task_name,
                            queue_name=queue_name,
                            task_event_queue=task_event_queue,
                            worker_id=worker_id,
                            worker_data=worker_data,
                            group_info_key=key
                        )

                        if success:
                            logger.debug(
                                f"[Recovery] Successfully transferred messages for "
                                f"worker={worker_id}, task={task_name}, queue={queue_name}"
                            )
                        else:
                            logger.debug(
                                f"[Recovery] Failed to transfer messages for "
                                f"worker={worker_id}, task={task_name}, queue={queue_name}"
                            )

                    except json.JSONDecodeError as e:
                        logger.error(f"[Recovery] Failed to parse group_info: {e}")
                    except Exception as e:
                        logger.error(f"[Recovery] Error processing group_info: {e}")

                logger.debug(f"[Recovery] Checking all group_info statuses for worker={worker_id}")

                fresh_worker_data = await self.get_worker_info(worker_id)
                if not fresh_worker_data:
                    logger.debug(f"[Recovery] Worker {worker_id} not found in Redis, skipping final check")
                    continue

                all_group_infos_transferred = True
                group_info_count = 0
                transferred_count = 0

                for key, value in fresh_worker_data.items():
                    if not key.startswith('group_info:'):
                        continue

                    group_info_count += 1
                    try:
                        group_info = json.loads(value)
                        is_transferred = group_info.get('messages_transferred', 'false')

                        if is_transferred == 'true':
                            transferred_count += 1
                            logger.debug(
                                f"[Recovery] Group_info {key} already transferred for worker={worker_id}"
                            )
                        else:
                            all_group_infos_transferred = False
                            logger.debug(
                                f"[Recovery] Group_info {key} NOT yet transferred for worker={worker_id}, "
                                f"cannot mark worker-level messages_transferred=true"
                            )
                    except json.JSONDecodeError as e:
                        logger.error(f"[Recovery] Failed to parse group_info {key}: {e}")
                        all_group_infos_transferred = False

                logger.debug(
                    f"[Recovery] Worker {worker_id} has {group_info_count} group_infos, "
                    f"{transferred_count} transferred"
                )

                if group_info_count == 0:
                    await self.async_redis.hset(
                        worker_key,
                        'messages_transferred',
                        'true'
                    )
                    logger.debug(
                        f"[Recovery] Worker {worker_id} has no group_infos, "
                        f"marked worker-level messages_transferred=true"
                    )
                elif all_group_infos_transferred:
                    await self.async_redis.hset(
                        worker_key,
                        'messages_transferred',
                        'true'
                    )
                    logger.debug(
                        f"[Recovery] All {group_info_count} group_infos transferred successfully for worker={worker_id}, "
                        f"marked worker-level messages_transferred=true"
                    )
                else:
                    logger.debug(
                        f"[Recovery] Not all group_infos transferred for worker={worker_id} "
                        f"({transferred_count}/{group_info_count}), will retry in next cycle"
                    )

            except Exception as e:
                logger.error(f"[Recovery] Error processing worker {worker_key}: {e}")
            finally:
                await worker_lock.release()
                logger.debug(f"[Recovery] Released lock for worker={worker_id}")

    async def _transfer_messages(
        self,
        stream_key: str,
        group_name: str,
        offline_consumer_name: str,
        task_name: str,
        queue_name: str,
        task_event_queue: asyncio.Queue,
        worker_id: str,
        worker_data: dict,
        group_info_key: str
    ) -> bool:

        total_transferred = 0

        try:
            messages_transferred = worker_data.get('messages_transferred', 'false')
            if isinstance(messages_transferred, bytes):
                messages_transferred = messages_transferred.decode('utf-8')

            if messages_transferred.lower() == 'true':
                logger.debug(f"[Recovery] Messages already transferred for worker={worker_id}")
                return False

            try:
                batch_size = 100

                while True:
                    detailed_pending = await self.async_redis.xpending_range(
                        stream_key,
                        group_name,
                        min='-',
                        max='+',
                        count=batch_size,
                        consumername=offline_consumer_name
                    )

                    if not detailed_pending:
                        logger.debug(
                            f"[Recovery] No pending messages for worker={worker_id}, "
                            f"task={task_name}, queue={queue_name}"
                        )
                        break

                    logger.debug(
                        f"[Recovery] Found {len(detailed_pending)} pending messages for "
                        f"worker={worker_id}, task={task_name}, queue={queue_name}"
                    )

                    message_ids = [msg['message_id'] for msg in detailed_pending]

                    claimed_messages = await self.async_redis.xclaim(
                        stream_key,
                        group_name,
                        self.worker_id,  
                        min_idle_time=0,
                        message_ids=message_ids
                    )

                    if not claimed_messages:
                        logger.debug(
                            f"[Recovery] Failed to claim messages for worker={worker_id}, "
                            f"task={task_name}, queue={queue_name}"
                        )
                        break

                    logger.debug(
                        f"[Recovery] Claimed {len(claimed_messages)} messages for "
                        f"worker={worker_id}, task={task_name}, queue={queue_name}"
                    )

                    for msg_id, msg_data in claimed_messages:
                        try:
                            if isinstance(msg_id, bytes):
                                msg_id = msg_id.decode('utf-8')

                            data_field = msg_data.get(b'data') or msg_data.get('data')
                            if not data_field:
                                logger.debug(f"[Recovery] No data field in message {msg_id}")
                                continue

                            parsed_data = msgpack.unpackb(data_field, raw=False)

                            parsed_data['_task_name'] = task_name
                            parsed_data['queue'] = queue_name

                            task_item = {
                                'queue': queue_name,
                                'event_id': msg_id,
                                'event_data': parsed_data,
                                'consumer': self.worker_id,
                                'group_name': group_name
                            }

                            await task_event_queue.put(task_item)
                            total_transferred += 1

                            logger.debug(
                                f"[Recovery] Put message {msg_id} into event_queue for task={task_name}"
                            )

                        except Exception as e:
                            logger.error(f"[Recovery] Error processing message {msg_id}: {e}")

            except Exception as e:
                logger.error(
                    f"[Recovery] Error during message recovery for worker={worker_id}, "
                    f"task={task_name}, queue={queue_name}: {e}, will still mark as transferred"
                )


            group_info_value = await self.async_redis.hget(
                self._get_worker_key(worker_id),
                group_info_key
            )

            if group_info_value:
                if isinstance(group_info_value, bytes):
                    group_info_value = group_info_value.decode('utf-8')

                group_info = json.loads(group_info_value)
                group_info['messages_transferred'] = 'true'

                await self.update_worker_field(
                    worker_id,
                    group_info_key,
                    json.dumps(group_info)
                )
                logger.debug(
                    f"[Recovery] Marked group_info transferred: worker={worker_id}, "
                    f"task={task_name}, queue={queue_name}, key={group_info_key}, transferred={total_transferred} messages"
                )
            else:
                logger.debug(
                    f"[Recovery] group_info not found in Redis for key={group_info_key}, "
                    f"worker={worker_id}, cannot mark as transferred"
                )

          

            return True  

        except Exception as e:
            logger.error(
                f"[Recovery] Error transferring messages for worker={worker_id}, "
                f"task={task_name}, queue={queue_name}: {e}"
            )
            return False  

    def stop_recovery(self):
        self._stop_recovery = True


    async def start_cleanup_task(self, cleanup_interval: float = 86400.0):
        logger.debug(f"[Cleanup] Starting offline worker cleanup task (interval={cleanup_interval}s)")
        self._stop_cleanup = False

        while not self._stop_cleanup:
            try:
                await self._cleanup_offline_workers_once()

                await asyncio.sleep(cleanup_interval)

            except asyncio.CancelledError:
                logger.debug("[Cleanup] Cleanup task cancelled")
                break
            except Exception as e:
                logger.error(f"[Cleanup] Unexpected error: {e}", exc_info=True)
                await asyncio.sleep(cleanup_interval)

        logger.debug("[Cleanup] Offline worker cleanup task stopped")

    async def _cleanup_offline_workers_once(self):
        current_time = time.time()
        one_day_ago = current_time - 86400  
        cleanup_count = 0

        try:
            all_worker_ids = await self.get_all_workers()
            logger.debug(f"[Cleanup] Scanning {len(all_worker_ids)} workers for cleanup")

            for worker_id in all_worker_ids:
                if self._stop_cleanup:
                    break

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

                lock_key = f"{self.redis_prefix}:CLEANUP:WORKER_LOCK:{worker_id}"
                worker_lock = AsyncLock(
                    self.async_redis,
                    lock_key,
                    timeout=60,  
                    blocking=False  
                )

                if not await worker_lock.acquire():
                    logger.debug(f"[Cleanup] Worker {worker_id} is being processed by another worker, skipping")
                    continue

                try:
                    worker_info = await self.get_worker_info(worker_id)

                    if not worker_info:
                        logger.debug(
                            f"[Cleanup] Worker {worker_id} in registry but hash not found, "
                            f"cleaning up inconsistent data"
                        )
                        await self.delete_worker(worker_id)
                        cleanup_count += 1
                        continue

                    is_alive = worker_info.get('is_alive', 'false')
                    messages_transferred = worker_info.get('messages_transferred', 'false')
                    offline_time_str = worker_info.get('offline_time', '0')

                    is_alive_bool = (is_alive == 'true')
                    messages_transferred_bool = (messages_transferred == 'true')

                    try:
                        offline_time = float(offline_time_str)
                    except (ValueError, TypeError):
                        offline_time = 0

                    should_cleanup = (
                        not is_alive_bool  
                        and messages_transferred_bool  
                        and offline_time > 0  
                        and offline_time < one_day_ago  
                    )

                    if should_cleanup:
                        offline_duration = current_time - offline_time
                        offline_hours = offline_duration / 3600

                        logger.debug(
                            f"[Cleanup] Cleaning up worker {worker_id}: "
                            f"offline for {offline_hours:.1f} hours, "
                            f"messages_transferred={messages_transferred}"
                        )

                        await self.delete_worker(worker_id)

                        cleanup_count += 1
                        logger.debug(f"[Cleanup] Successfully cleaned up worker {worker_id}")

                    else:
                        if is_alive_bool:
                            reason = "still alive"
                        elif not messages_transferred_bool:
                            reason = "messages not transferred"
                        elif offline_time == 0:
                            reason = "no offline_time"
                        elif offline_time >= one_day_ago:
                            offline_duration = current_time - offline_time
                            offline_hours = offline_duration / 3600
                            reason = f"offline for only {offline_hours:.1f} hours (< 24h)"
                        else:
                            reason = "unknown"

                        logger.debug(
                            f"[Cleanup] Worker {worker_id} not eligible for cleanup: {reason}"
                        )

                except Exception as e:
                    logger.error(f"[Cleanup] Error processing worker {worker_id}: {e}", exc_info=True)
                finally:
                    await worker_lock.release()

            if cleanup_count > 0:
                logger.debug(f"[Cleanup] Cleaned up {cleanup_count} offline workers")
            else:
                logger.debug("[Cleanup] No workers eligible for cleanup")

        except Exception as e:
            logger.error(f"[Cleanup] Error during cleanup: {e}", exc_info=True)

    def stop_cleanup(self):
        self._stop_cleanup = True


__all__ = [
    'WorkerManager',
]
