import json
import sys

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

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('____', 1)
        except redis_pkg.exceptions.ConnectionError as e:
            print('Failed to connect to the Redis server:', e)
            print('Did you install and start the Redis server?')
            sys.exit(1)

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

    def set(self, key, value):
        return self.redis.set(self._key(key), value)

    def get(self, key):
        return self.redis.get(self._key(key))

    def incr(self, key, amount=1):
        return self.redis.incr(self._key(key), amount)

    def decr(self, key, amount=1):
        return self.redis.decr(self._key(key), amount)

    def sadd(self, key: str, *values: FieldT):
        return self.redis.sadd(self._key(key), *values)

    def srem(self, key: str, *values: FieldT):
        return self.redis.srem(self._key(key), *values)

    def sismember(self, key: str, value: str):
        return self.redis.sismember(self._key(key), value)

    def set_dict(self, key, dict_value):
        return self.set(self._key(key), json.dumps(dict_value))

    def get_dict(self, key):
        r = self.get(self._key(key))
        if not r:
            return dict()
        else:
            return json.loads(r.decode("utf-8"))

    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')