import redis
from django.conf import settings
from django.core.management import BaseCommand

'''
日记相关redis 部分和gaia 分隔
'''
class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument('-t', '--test', action='store_true',
                            help="Run test for redis migrate, a redis server should run on localhost:6379")
        parser.add_argument('-m', '--migrate', action='store_true', help="Redis data migrate")

    def handle(self, *args, **options):
        """
        redis migrate.
        :param args:
        :param options:
        :return:
        """
        for rd in settings.REDIS_MIGRATE_TALOS:
            for prefix in rd['prefix']:
                self.migrate_v2(rd['source'], rd['target'], prefix=prefix)


    def migrate_v2(self, source, target, prefix=""):
        """
        适用于：redis >= 2.6
        :param source:
        :param target:
        :param prefix:
        :return:
        """
        self.stdout.write(self.style.SUCCESS(
            "Start to migrate from %s to %s, prefix: %s" % (
                "redis://" + source['host'] + ":" + str(source['port']) + "/" + str(source['db']),
                "redis://" + target['host'] + ":" + str(target['port']) + "/" + str(target['db']),
                prefix
            )
        ))
        src_rd = redis.StrictRedis(**source)
        dst_rd = redis.StrictRedis(**target)
        total = len(src_rd.keys(prefix + "*"))
        for key in src_rd.keys(prefix + "*"):
            try:
                """ redis> SET """
                v = src_rd.get(key)
                dst_rd.set(key, v)
            except redis.exceptions.ResponseError:
                """ redis> HSET """
                try:
                    v = src_rd.hgetall(key)
                    dst_rd.hmset(key, v)
                except redis.exceptions.ResponseError:
                    """redis> ZADD """
                    zv = src_rd.zrange(key, 0, -1, withscores=True)
                    v = {key.decode(): value for key, value in zv}
                    dst_rd.zadd(key, **v)

        self.stdout.write(self.style.SUCCESS("Migrate success: %s key." % total))