import json
import sys

from flask_caching import Cache
from redis import Redis
from redis.typing import FieldT
import redis as redis_pkg

cache = Cache(config={'CACHE_TYPE': 'RedisCache', 'CACHE_REDIS_URL': 'redis://localhost:6379/0', 'CACHE_KEY_PREFIX': 'local-llm'})


# redis = Redis()

class RedisWrapper:
    """
    A wrapper class to set prefixes to keys.
    """

    def __init__(self, prefix, **kwargs):
        self.redis = Redis(**kwargs)
        self.prefix = prefix
        try:
            self.set('check_connected', 1)
        except redis_pkg.exceptions.ConnectionError as e:
            print('Failed to connect to the Redis server:', e)
            print('Did you install and start Redis?')
            sys.exit(1)

    def set(self, key, value):
        return self.redis.set(f"{self.prefix}:{key}", value)

    def get(self, key):
        return self.redis.get(f"{self.prefix}:{key}")

    def incr(self, key, amount=1):
        return self.redis.incr(f"{self.prefix}:{key}", amount)

    def decr(self, key, amount=1):
        return self.redis.decr(f"{self.prefix}:{key}", amount)

    def sadd(self, key: str, *values: FieldT):
        return self.redis.sadd(f"{self.prefix}:{key}", *values)

    def srem(self, key: str, *values: FieldT):
        return self.redis.srem(f"{self.prefix}:{key}", *values)

    def sismember(self, key: str, value: str):
        return self.redis.sismember(f"{self.prefix}:{key}", value)

    def set_dict(self, key, dict_value):
        # return self.redis.hset(f"{self.prefix}:{key}", mapping=dict_value)
        return self.set(f"{self.prefix}:{key}", json.dumps(dict_value))

    def get_dict(self, key):
        # return self.redis.hgetall(f"{self.prefix}:{key}")
        r = self.get(f"{self.prefix}:{key}")
        if not r:
            return dict()
        else:
            return json.loads(r)

    def flush(self):
        flushed = []
        for key in self.redis.scan_iter(f'{self.prefix}:*'):
            flushed.append(key)
            self.redis.delete(key)
        return flushed


redis = RedisWrapper('local_llm')