from typing import Dict, Optional

from redisaq.utils import (
    get_redis_key_for_topic_consumer_group,
    get_redis_key_for_topic_consumer_group_consumer,
    get_redis_key_for_topic_partition,
    get_redis_key_for_topic_partition_messages,
    get_redis_key_for_topic_rebalance_channel,
)


class TopicConsumerGroupKeys:
    def __init__(self, topic: str, consumer_group: str):
        self.group_name = consumer_group
        self.rebalance_channel = get_redis_key_for_topic_rebalance_channel(
            topic=topic, consumer_group=consumer_group
        )
        self.consumer_key = get_redis_key_for_topic_consumer_group_consumer(
            topic=topic, consumer_group=consumer_group
        )


class TopicPartitionKeys:
    def __init__(self, topic: str, partition: int):
        self.stream_key = get_redis_key_for_topic_partition_messages(
            topic=topic, partition=partition
        )


class TopicKeys:
    def __init__(self, topic: str):
        self.topic = topic
        self.partition_key = get_redis_key_for_topic_partition(topic)
        self.consumer_group_key = get_redis_key_for_topic_consumer_group(topic)
        self.partition_keys: Dict[int, TopicPartitionKeys] = {}
        self.consumer_group_keys = TopicConsumerGroupKeys(self.topic, "default")

    def has_partition(self, partition: int) -> bool:
        return partition in self.partition_keys

    def add_partition(self, partition: int) -> None:
        self.partition_keys[partition] = TopicPartitionKeys(
            topic=self.topic, partition=partition
        )

    def set_consumer_group(self, consumer_group: str) -> None:
        self.consumer_group_keys = TopicConsumerGroupKeys(
            topic=self.topic,
            consumer_group=consumer_group,
        )
