"""
队列和任务注册管理模块
负责队列、任务、消费者组的注册和查询功能

整合了原来分散的注册功能：
- 队列注册 (REGISTRY:QUEUES)
- 任务注册 (REGISTRY:TASKS)
- 消费者组注册 (REGISTRY:CONSUMER_GROUPS)
"""

import logging
from typing import Set, List, Dict, Any

logger = logging.getLogger(__name__)


class QueueRegistry:
    """
    队列和任务注册管理器

    维护队列和任务的注册信息，提供发现功能。
    整合了原来 TaskRegistry 的 Redis 注册功能。

    Redis 键结构：
    - {prefix}:REGISTRY:QUEUES - 所有队列（包括优先级队列）
    - {prefix}:REGISTRY:TASKS - 所有任务名称
    - {prefix}:REGISTRY:CONSUMER_GROUPS:{queue} - 队列的消费者组
    """

    def __init__(self, redis_client, async_redis_client, redis_prefix: str = 'jettask'):
        self.redis = redis_client

        self.async_redis = async_redis_client

        self.redis_prefix = redis_prefix

        self.queues_registry_key = f"{redis_prefix}:REGISTRY:QUEUES"  
        self.tasks_registry_key = f"{redis_prefix}:REGISTRY:TASKS"  
        self.consumer_groups_registry_key = f"{redis_prefix}:REGISTRY:CONSUMER_GROUPS"


    async def register_queue(self, queue_name: str):
        await self.async_redis.sadd(self.queues_registry_key, queue_name)
        logger.debug(f"Registered queue: {queue_name}")

    async def unregister_queue(self, queue_name: str):
        await self.async_redis.srem(self.queues_registry_key, queue_name)
        logger.debug(f"Unregistered queue: {queue_name}")

    async def get_all_queues(self) -> Set[str]:
        return await self.async_redis.smembers(self.queues_registry_key)

    async def get_queue_count(self) -> int:
        return await self.async_redis.scard(self.queues_registry_key)

    async def get_base_queues(self) -> Set[str]:
        all_queues = await self.get_all_queues()

        base_queues = set()
        for queue in all_queues:
            if isinstance(queue, bytes):
                queue = queue.decode('utf-8')

            parts = queue.split(':')
            if len(parts) >= 2 and parts[-1].isdigit():
                base_queue = ':'.join(parts[:-1])
                base_queues.add(base_queue)
            else:
                base_queues.add(queue)

        return base_queues

    async def discover_matching_queues(self, wildcard_pattern: str) -> Set[str]:
        from jettask.utils.queue_matcher import discover_matching_queues

        all_registered_queues = await self.get_all_queues()

        all_registered_queues = {
            q.decode('utf-8') if isinstance(q, bytes) else q
            for q in all_registered_queues
        }

        matched_queues = discover_matching_queues([wildcard_pattern], all_registered_queues)
        return matched_queues


    async def register_consumer_group(self, queue: str, group_name: str):
        key = f"{self.consumer_groups_registry_key}:{queue}"
        await self.async_redis.sadd(key, group_name)
        logger.debug(f"Registered consumer group: {group_name} for queue: {queue}")

    async def unregister_consumer_group(self, queue: str, group_name: str):
        key = f"{self.consumer_groups_registry_key}:{queue}"
        await self.async_redis.srem(key, group_name)
        logger.debug(f"Unregistered consumer group: {group_name} for queue: {queue}")

    async def get_consumer_groups_for_queue(self, queue: str) -> Set[str]:
        key = f"{self.consumer_groups_registry_key}:{queue}"
        return await self.async_redis.smembers(key)


    async def register_priority_queue(self, base_queue: str, priority: int):
        priority_queue = f"{base_queue}:{priority}"
        await self.async_redis.sadd(self.queues_registry_key, priority_queue)
        logger.debug(f"Registered priority queue: {priority_queue}")

    async def unregister_priority_queue(self, base_queue: str, priority: int):
        priority_queue = f"{base_queue}:{priority}"
        await self.async_redis.srem(self.queues_registry_key, priority_queue)
        logger.debug(f"Unregistered priority queue: {priority_queue}")

    async def get_priority_queues_for_base(self, base_queue: str) -> List[str]:
        all_queues = await self.async_redis.smembers(self.queues_registry_key)

        result = []
        for queue in all_queues:
            if isinstance(queue, bytes):
                queue = queue.decode('utf-8')

            if queue.startswith(f"{base_queue}:"):
                parts = queue.split(':')
                if len(parts) >= 2 and parts[-1].isdigit():
                    result.append(queue)

        result.sort(key=lambda x: int(x.split(':')[-1]))
        return result

    async def clear_priority_queues_for_base(self, base_queue: str):
        priority_queues = await self.get_priority_queues_for_base(base_queue)

        if priority_queues:
            await self.async_redis.srem(self.queues_registry_key, *priority_queues)
            logger.debug(f"Cleared {len(priority_queues)} priority queues for base queue: {base_queue}")


    async def get_task_names_by_queue(self, base_queue: str) -> Set[str]:
        read_offsets_key = f"{self.redis_prefix}:READ_OFFSETS"

        all_keys = await self.async_redis.hkeys(read_offsets_key)

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

            if not key.startswith(f"{base_queue}:"):
                continue

            suffix = key[len(base_queue) + 1:]  

            parts = suffix.split(':')

            if parts[0].isdigit() and len(parts) > 1:
                task_name = ':'.join(parts[1:])
            else:
                task_name = suffix

            if task_name:  
                task_names.add(task_name)

        return task_names


    async def register_task(self, task_name: str, queue: str):
        if not self.async_redis:
            logger.warning("Redis client not configured, skipping task registration")
            return

        await self.async_redis.hset(self.tasks_registry_key, task_name, queue)
        logger.debug(f"Registered task: {task_name} -> {queue}")

    async def unregister_task(self, task_name: str):
        if not self.async_redis:
            return

        await self.async_redis.hdel(self.tasks_registry_key, task_name)
        logger.debug(f"Unregistered task: {task_name}")

    async def get_all_tasks(self) -> Set[str]:
        if not self.async_redis:
            return set()

        tasks = await self.async_redis.hkeys(self.tasks_registry_key)
        return {
            t.decode('utf-8') if isinstance(t, bytes) else t
            for t in tasks
        }

    async def get_task_count(self) -> int:
        if not self.async_redis:
            return 0

        return await self.async_redis.hlen(self.tasks_registry_key)

    async def task_exists(self, task_name: str) -> bool:
        if not self.async_redis:
            return False

        return await self.async_redis.hexists(self.tasks_registry_key, task_name)

    async def get_queue_for_task(self, task_name: str) -> str:
        if not self.async_redis:
            return None

        queue = await self.async_redis.hget(self.tasks_registry_key, task_name)
        if queue and isinstance(queue, bytes):
            queue = queue.decode('utf-8')
        return queue

    async def get_all_task_queue_pairs(self) -> List[Dict[str, str]]:
        if not self.async_redis:
            return []

        task_queue_map = await self.async_redis.hgetall(self.tasks_registry_key)

        pairs = []
        for task_name, queue in task_queue_map.items():
            if isinstance(task_name, bytes):
                task_name = task_name.decode('utf-8')
            if isinstance(queue, bytes):
                queue = queue.decode('utf-8')

            pairs.append({
                'queue': queue,
                'task_name': task_name
            })

        return pairs

    async def get_tasks_by_queue(self, queue: str) -> List[str]:
        if not self.async_redis:
            return []

        base_queue = queue
        parts = queue.split(':')
        if len(parts) >= 2 and parts[-1].isdigit():
            base_queue = ':'.join(parts[:-1])

        task_queue_map = await self.async_redis.hgetall(self.tasks_registry_key)

        tasks = []
        for task_name, task_queue in task_queue_map.items():
            if isinstance(task_name, bytes):
                task_name = task_name.decode('utf-8')
            if isinstance(task_queue, bytes):
                task_queue = task_queue.decode('utf-8')

            task_base_queue = task_queue
            task_parts = task_queue.split(':')
            if len(task_parts) >= 2 and task_parts[-1].isdigit():
                task_base_queue = ':'.join(task_parts[:-1])

            if task_base_queue == base_queue or task_queue == queue:
                tasks.append(task_name)

        return tasks
