import dataclasses
import functools
import itertools
import logging
import re
import time
import typing
import uuid
from typing import Optional

import boto3
import botocore.exceptions
from botocore.exceptions import ClientError

logger = logging.getLogger(__name__)


def retry(timeout: int, backoff_seconds: int = 10, bubble_errors: typing.List[typing.Type[Exception]] = None):
    """
    Tries calling a function {timeout} seconds until it succeeds or gives up and throw an error
    :param timeout: Timeout in seconds
    :param backoff_seconds: Wait time between retries
    :param bubble_errors: List of exceptions to bubble up
    :return: Wrapped function
    """
    bubble_errors = tuple(bubble_errors or [])

    def wrapped(func: typing.Callable):
        @functools.wraps(func)
        def inner(*args, **kwargs):
            waited = 0
            while True:
                try:
                    logger.debug(f"Calling {func}")
                    return func(*args, **kwargs)
                except KeyboardInterrupt:
                    raise
                except tuple(bubble_errors):
                    raise
                except Exception as e:
                    doze = min(timeout - waited, backoff_seconds)
                    times = (timeout - waited) // backoff_seconds
                    logger.debug(f"{func} failed, will try {times} times in {doze} sec")
                    time.sleep(doze)
                    waited += backoff_seconds
                    if waited >= timeout:
                        raise TimeoutError(f"Couldn't get a successful result from {func} in {timeout} seconds") from e

        return inner

    return wrapped


@dataclasses.dataclass
class RedisInstance:
    cluster_id: str
    host: str
    port: int
    arn: str
    user_group_ids: typing.List[str] = dataclasses.field(default_factory=list)


def paginate_by_marker(
    func: callable,
    list_field: str,
) -> typing.Iterable:
    """
    Simplifies paginating over AWS results with a simpler interface than Paginators.
    :param func: A function that accepts a parameter named `Marker`
    :param list_field: Result field to iterate on
    """
    marker = ""
    while True:
        result = func(Marker=marker)

        items = result.get(list_field, [])
        yield from items

        if "Marker" not in result:
            break

        marker = result["Marker"]


@dataclasses.dataclass
class RedisUser:
    username: str
    password: str
    redis_key_prefix: str


class CreateUserResult(typing.TypedDict):
    user_group_id: str
    users: typing.List[RedisUser]
    status: str


class RedisService:
    def __init__(self, region: str):
        self.elasticache = boto3.client("elasticache", region_name=region)

    def create_redis_acl_user(
        self,
        replication_group_id: str,
        users: typing.List[RedisUser],
    ) -> None:
        """
        Creates a Redis ACL user. It will throw an error if the user group attached to the replication group is not ready.
        This prevents us from quickly creating users. In that case, pass in a list of users instead.

        It takes about 2 minutes until the user is ready.
        """
        repl_groups = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)
        user_group_id: str = repl_groups['ReplicationGroups'][0]['UserGroupIds'][0]
        group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)['UserGroups'][0]

        """
        Limits imposed by AWS:
        
        - User groups per replication group = 1
        - Users per user group = 100
        - Number of users = 1000
        - Number of user groups = 100
        https://docs.aws.amazon.com/AmazonElastiCache/latest/red-ug/Clusters.RBAC.html#Users-groups-to-RGs
        
        In short, we can't place more than 100 users on a single Redis instance.   
        """

        total_users = len(group['UserIds'])
        if total_users >= 100:
            raise Exception('Replication group reached the limit of 100 users')

        self._update_or_create_user_group(user_group_id, users)

    def delete_redis_acl_user(self, username: str) -> None:
        logger.info(f'Deleting user {username}')
        self.elasticache.delete_user(UserId=username)

        # Even though we just deleted the user, it takes a couple of seconds until AWS updates user group status.
        # So, we have to wait a bit before we can actually wait and check that the changes have propagated.
        # Overall, it should take about 2-3 minutes to delete a user.

    def create_redis_acl(
        self,
        *,
        replication_group_id: str,
        description: str = None,
        node_type: str,
        cache_subnet_group_name: str,
        security_group_id: str,
        tags: typing.Dict[str, str] = None,
        users: typing.List[RedisUser],
    ) -> dict:
        user_group_id = f"{replication_group_id}-ug"
        # we have to wait (~1m) until the user group is ready before we can create the ACL
        self._update_or_create_user_group(user_group_id, users, wait=True)

        cache_param_group = 'redis-acl-with-100-db'
        self._ensure_cache_parameter_group(cache_param_group)

        logger.info(f"Creating Redis replication group {replication_group_id=}")
        default_replication_group_kwargs = dict(
            ReplicasPerNodeGroup=0,
            Engine="redis",
            EngineVersion="6.x",
            TransitEncryptionEnabled=True,  # must be True to use Redis ACL
            MultiAZEnabled=False,
            AutomaticFailoverEnabled=False,
        )
        repl_result = self.elasticache.create_replication_group(
            **default_replication_group_kwargs,
            ReplicationGroupId=replication_group_id,
            ReplicationGroupDescription=description or f'{replication_group_id} replication group',
            CacheNodeType=node_type,
            CacheSubnetGroupName=cache_subnet_group_name,
            CacheParameterGroupName=cache_param_group,
            UserGroupIds=[
                user_group_id,
            ],
            SecurityGroupIds=[
                security_group_id,
            ],
            Tags=[{"Key": k, "Value": v} for k, v in (tags or {}).items()],
        )

        # replication group takes ~8m minutes to be ready

        return {
            "replication_group_id": repl_result["ReplicationGroup"]["ReplicationGroupId"],
            'arn': repl_result["ReplicationGroup"]["ARN"],
            'status': repl_result["ReplicationGroup"]["Status"],
            'users': users,
            '_response': repl_result,
        }

    def _update_or_create_user_group(
        self, user_group_id: str, users: typing.List[RedisUser], wait: bool = False
    ) -> None:
        """
        Creates or updates a user group with the given users. Waits until the user group is active.

        :param user_group_id:
        :param users:
        """
        # we have to add the `default` user for backwards compatibility
        existing_users = self.elasticache.describe_users(
            Filters=[{"Name": "user-id", "Values": [u.username for u in users]}],
        )["Users"]
        existing_user_ids = [it["UserId"] for it in existing_users]

        created_user_ids = ["default"]
        for user in users:
            if user.username in existing_user_ids:
                logger.info(f"User {user.username} is already exists")
                created_user_ids.append(user.username)
                continue

            redis_acl = f"on +@all -@dangerous ~{user.redis_key_prefix}*"
            logger.info(f"Creating tenant {user.username=}")
            user_result = self.elasticache.create_user(
                UserId=user.username,
                UserName=user.username,
                Passwords=[user.password],
                AccessString=redis_acl,
                Engine="redis",
            )
            created_user_ids.append(user_result["UserId"])

        try:
            group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)["UserGroups"][0]
            member_user_ids: typing.List[str] = group["UserIds"]
            users_to_add = set(created_user_ids) - set(member_user_ids)

            if users_to_add:
                self.elasticache.modify_user_group(
                    UserGroupId=user_group_id,
                    UserIdsToAdd=list(users_to_add),
                )
                # this will take some time (~45s) until the changes propagate
        except self.elasticache.exceptions.UserGroupNotFoundFault:
            logger.info('User group does not exist, creating')
            _ = self.elasticache.create_user_group(
                UserGroupId=user_group_id,
                Engine="redis",
                UserIds=created_user_ids,
            )
            # user group creation takes around 60s
            if wait:
                self._wait_user_group(user_group_id)

    @retry(timeout=60, backoff_seconds=5)
    def _wait_user_group(self, user_group_id: str) -> None:
        logger.info(f'Checking status of user group {user_group_id=}')
        g = self.elasticache.describe_user_groups(UserGroupId=user_group_id)["UserGroups"][0]
        assert g["Status"] == "active"

    def is_redis_acl_ready(self, replication_group_id: str) -> bool:
        repl = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)[
            'ReplicationGroups'
        ][0]
        is_redis_available = repl['Status'] == 'available'

        user_group_id = repl['UserGroupIds'][0]
        group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)['UserGroups'][0]
        is_group_available = group['Status'] == 'active'

        return is_group_available and is_redis_available

    def get_redis_acl(self, replication_group_id: str) -> dict:
        redis = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)[
            'ReplicationGroups'
        ][0]
        return dict(
            replication_group_id=redis['ReplicationGroupId'],
            arn=redis['ARN'],
            host=redis["NodeGroups"][0]["PrimaryEndpoint"]["Address"],
            port=redis["NodeGroups"][0]["PrimaryEndpoint"]["Port"],
        )

    def delete_redis_acl(self, replication_group_id: str) -> None:
        self.elasticache.delete_replication_group(ReplicationGroupId=replication_group_id)

    def _ensure_cache_parameter_group(self, parameter_group_name: str) -> None:
        try:
            _ = self.elasticache.describe_cache_parameter_groups(CacheParameterGroupName=parameter_group_name)[
                'CacheParameterGroups'
            ][0]
            return
        except self.elasticache.exceptions.CacheParameterGroupNotFoundFault:
            pass

        logger.info(f"Creating cache parameter group {parameter_group_name=}")
        _ = self.elasticache.create_cache_parameter_group(
            CacheParameterGroupName=parameter_group_name,
            CacheParameterGroupFamily="redis6.x",
            Description="Redis parameter group for the Redis ACL",
        )
        _ = self.elasticache.modify_cache_parameter_group(
            CacheParameterGroupName=parameter_group_name,
            ParameterNameValues=[
                {
                    "ParameterName": "databases",
                    "ParameterValue": "100",
                },
            ],
        )


def akinon_create_redis(
    *,
    k8s_cluster_name: str,
    region: str,
    node_type: str,
    replicas: int,
    owner_arn: str,
    app_name: str,
    role: str,
):
    subnet_group_name = f"{k8s_cluster_name}-redis-subg"
    security_group_name = f"{k8s_cluster_name}-redis-sg"

    # owner arn is formatted like:
    # arn:aws:iam::412344683105:user/myusername
    username = owner_arn.split("/")[1]
    tags = {
        "CostCenter": f"{username}-{app_name}-{role}-redis",
    }

    cache_cluster_id = uuid.uuid4().hex
    r = RedisService(region)
    return r.create_redis(
        cluster_id=cache_cluster_id,
        node_type=node_type,
        replicas=replicas,
        security_group_name=security_group_name,
        subnet_group_name=subnet_group_name,
        tags=tags,
    )


def akinon_create_redis_with_acl(
    k8s_cluster_name: str,
    owner_arn: str,
    node_type: str,
    app_name: str,
    role: str,
    users: typing.List[RedisUser],
) -> RedisInstance:
    cache_cluster_id = uuid.uuid4().hex
    subnet_group_name = f"{k8s_cluster_name}-redis-subg"
    security_group_name = f"{k8s_cluster_name}-redis-sg"

    # owner arn is formatted as:
    # arn:aws:iam::412344683105:user/myusername
    username = owner_arn.split("/")[1]
    tags = {
        "CostCenter": f"{username}-{app_name}-{role}-redis",
    }

    return r.create_redis_acl(
        cluster_id=cache_cluster_id,
        node_type=node_type,
        subnet_group_name=subnet_group_name,
        security_group_name=security_group_name,
        tags=tags,
        users=users,
    )


def chunk(it, size):
    it = iter(it)
    sentinel = ()
    return iter(lambda: tuple(itertools.islice(it, size)), sentinel)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format=f"%(asctime)s: {logging.BASIC_FORMAT}")
    r = RedisService(region="eu-central-1")
    # res = r.create_redis_acl(
    #     replication_group_id="zerodev3",
    #     description="zero redis dev",
    #     node_type="cache.t3.micro",
    #     cache_subnet_group_name="dev",
    #     security_group_id="sg-c11a77ac",
    #     tags={"CostCenter": "zero123"},
    #     users=[RedisUser(username="u1", password="zeropassword1234", redis_key_prefix="zero:")],
    # )
    # print(res)

    r.create_redis_acl_user(
        replication_group_id='zerodev3',
        users=[
            RedisUser(
                username='u2',
                password='testtesttesttest123',
                redis_key_prefix='test:',
            )
        ],
    )
    # r.delete_redis_acl_user('test6')
    while True:
        if r.is_redis_acl_ready('zerodev3'):
            logger.info('ready')
            print(r.get_redis_acl('zerodev3'))
            break
        logger.info('not ready')
        time.sleep(10)
    exit()

    res = r.create_redis_acl(
        replication_group_id="zerodev",
        description="zero redis dev",
        node_type="cache.t3.micro",
        cache_subnet_group_name="dev",
        security_group_id="sg-c11a77ac",
        tags={"CostCenter": "zero123"},
        users=[RedisUser(username="zerouser", password="zeropassword1234", redis_key_prefix="zero:")],
    )
    print(res)
    #
    # redis = r.get_redis_acl(cluster_arn="arn:aws:elasticache:eu-central-1:400344683105:replicationgroup:zero")
    # r.add_redis_user(
    #     cluster_id=redis.cluster_id,
    #     users=[RedisUser(username="zero27", password="zero27password123", redis_key_prefix="zero27:")],
    # )
    # print(r.add_redis_user(cluster_id="zero1", users=[RedisUser("zero26", "zero26password123", "zero26:")]))
    #
    # print(
    #     r.create_redis_with_acl(
    #         cluster_id="zero",
    #         cluster_description="zero redis",
    #         node_type="cache.t3.micro",
    #         subnet_group_name="testing-subnet",
    #         security_group_name="default",
    #         tags={"CostCenter": "zero123"},
    #         users=[RedisUser(username="zerouser", password="zeropassword1234", redis_key_prefix="zero:")],
    #     )
    # )