# Source code for rl_coach.memories.backend.redis

    #
    # Copyright (c) 2017 Intel Corporation
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    #

    import redis
    import pickle
    import uuid
    import time

    from rl_coach.memories.backend.memory import MemoryBackend, MemoryBackendParameters
    from rl_coach.core_types import Transition, Episode, EnvironmentSteps, EnvironmentEpisodes

    class RedisPubSubMemoryBackendParameters(MemoryBackendParameters):

        def __init__(self, redis_address: str="", redis_port: int=6379, channel: str="channel-{}".format(uuid.uuid4()),
                     orchestrator_params: dict=None, run_type='trainer', orchestrator_type: str = "kubernetes", deployed: str = False):
            self.redis_address = redis_address
            self.redis_port = redis_port
            self.channel = channel
            if not orchestrator_params:
                orchestrator_params = {}
            self.orchestrator_params = orchestrator_params
            self.run_type = run_type
            self.store_type = "redispubsub"
            self.orchestrator_type = orchestrator_type
            self.deployed = deployed

    class RedisPubSubBackend(MemoryBackend):
        """
        A memory backend which transfers the experiences from the rollout to the training worker using Redis Pub/Sub in
        Coach when distributed mode is used.
        """

        def __init__(self, params: RedisPubSubMemoryBackendParameters):
            """
            :param params: The Redis parameters to be used with this Redis Pub/Sub instance.
            """
            self.params = params
            self.redis_connection = redis.Redis(self.params.redis_address, self.params.redis_port)
            self.redis_server_name = 'redis-server-{}'.format(uuid.uuid4())
            self.redis_service_name = 'redis-service-{}'.format(uuid.uuid4())

        def store(self, obj):
            """
            :param obj: The object to store in memory. The object is either a Tranisition or Episode type.
            """
            self.redis_connection.publish(self.params.channel, pickle.dumps(obj))

        def deploy(self):
            """
            Deploy the Redis Pub/Sub service in an orchestrator.
            """
            if not self.params.deployed:
                if self.params.orchestrator_type == 'kubernetes':
                    self.deploy_kubernetes()

            # Wait till subscribe to the channel is possible or else it will cause delays in the trainer.
            time.sleep(10)

        def deploy_kubernetes(self):
            """
            Deploy the Redis Pub/Sub service in Kubernetes orchestrator.
            """
            if 'namespace' not in self.params.orchestrator_params:
                self.params.orchestrator_params['namespace'] = "default"
            from kubernetes import client

            container = client.V1Container(
                name=self.redis_server_name,
                image='redis:4-alpine',
            )
            template = client.V1PodTemplateSpec(
                metadata=client.V1ObjectMeta(labels={'app': self.redis_server_name}),
                spec=client.V1PodSpec(
                    containers=[container]
                )
            )
            deployment_spec = client.V1DeploymentSpec(
                replicas=1,
                template=template,
                selector=client.V1LabelSelector(
                    match_labels={'app': self.redis_server_name}
                )
            )

            deployment = client.V1Deployment(
                api_version='apps/v1',
                kind='Deployment',
                metadata=client.V1ObjectMeta(name=self.redis_server_name, labels={'app': self.redis_server_name}),
                spec=deployment_spec
            )

            api_client = client.AppsV1Api()
            try:
                api_client.create_namespaced_deployment(self.params.orchestrator_params['namespace'], deployment)
            except client.rest.ApiException as e:
                print("Got exception: %s\n while creating redis-server", e)
                return False

            core_v1_api = client.CoreV1Api()

            service = client.V1Service(
                api_version='v1',
                kind='Service',
                metadata=client.V1ObjectMeta(
                    name=self.redis_service_name
                ),
                spec=client.V1ServiceSpec(
                    selector={'app': self.redis_server_name},
                    ports=[client.V1ServicePort(
                        protocol='TCP',
                        port=6379,
                        target_port=6379
                    )]
                )
            )

            try:
                core_v1_api.create_namespaced_service(self.params.orchestrator_params['namespace'], service)
                self.params.redis_address = '{}.{}.svc'.format(
                    self.redis_service_name, self.params.orchestrator_params['namespace']
                )
                self.params.redis_port = 6379
                return True
            except client.rest.ApiException as e:
                print("Got exception: %s\n while creating a service for redis-server", e)
                return False

        def undeploy(self):
            """
            Undeploy the Redis Pub/Sub service in an orchestrator.
            """
            from kubernetes import client
            if self.params.deployed:
                return

            from kubernetes import client
            api_client = client.AppsV1Api()
            delete_options = client.V1DeleteOptions()
            try:
                api_client.delete_namespaced_deployment(self.redis_server_name, self.params.orchestrator_params['namespace'], delete_options)
            except client.rest.ApiException as e:
                print("Got exception: %s\n while deleting redis-server", e)

            api_client = client.CoreV1Api()
            try:
                api_client.delete_namespaced_service(self.redis_service_name, self.params.orchestrator_params['namespace'], delete_options)
            except client.rest.ApiException as e:
                print("Got exception: %s\n while deleting redis-server", e)

        def sample(self, size):
            pass

        def fetch(self, num_consecutive_playing_steps=None):
            """
            :param num_consecutive_playing_steps: The number steps to fetch.
            """
            return RedisSub(redis_address=self.params.redis_address, redis_port=self.params.redis_port, channel=self.params.channel).run(num_consecutive_playing_steps)

        def subscribe(self, agent):
            """
            :param agent: The agent in use.
            """
            redis_sub = RedisSub(redis_address=self.params.redis_address, redis_port=self.params.redis_port, channel=self.params.channel)
            return redis_sub

        def get_endpoint(self):
            return {'redis_address': self.params.redis_address,
                    'redis_port': self.params.redis_port}

    class RedisSub(object):
        def __init__(self, redis_address: str = "localhost", redis_port: int=6379, channel: str = "PubsubChannel"):
            super().__init__()
            self.redis_connection = redis.Redis(redis_address, redis_port)
            self.pubsub = self.redis_connection.pubsub()
            self.subscriber = None
            self.channel = channel
            self.subscriber = self.pubsub.subscribe(self.channel)

        def run(self, num_consecutive_playing_steps):
            """
            :param num_consecutive_playing_steps: The number steps to fetch.
            """
            transitions = 0
            episodes = 0
            steps = 0
            for message in self.pubsub.listen():
                if message and 'data' in message:
                    try:
                        obj = pickle.loads(message['data'])
                        if type(obj) == Transition:
                            transitions += 1
                            if obj.game_over:
                                episodes += 1
                            yield obj
                        elif type(obj) == Episode:
                            episodes += 1
                            transitions += len(obj.transitions)
                            yield from obj.transitions
                    except Exception:
                        continue

                if type(num_consecutive_playing_steps) == EnvironmentSteps:
                    steps = transitions
                if type(num_consecutive_playing_steps) == EnvironmentEpisodes:
                    steps = episodes

                if steps >= num_consecutive_playing_steps.num_steps:
                    break