import hashlib
import pickle

from llm_server.custom_redis import RedisCustom


class RedisClusterStore:
    def __init__(self, name: str, **kwargs):
        self.name = name
        self.config_redis = RedisCustom(name, **kwargs)

    def clear(self):
        self.config_redis.flush()

    def load(self, config: dict):
        for k, v in config.items():
            self.add_backend(k, v)

    def add_backend(self, name: str, values: dict):
        self.config_redis.hset(name, mapping={k: pickle.dumps(v) for k, v in values.items()})
        self.set_backend_value(name, 'online', False)
        h = hashlib.sha256(name.encode('utf-8')).hexdigest()
        self.set_backend_value(name, 'hash', f'{h[:8]}-{h[-8:]}')

    def set_backend_value(self, backend: str, key: str, value):
        # By storing the value as a pickle we don't have to cast anything when getting the value from Redis.
        self.config_redis.hset(backend, key, pickle.dumps(value))

    def get_backend(self, name: str):
        r = self.config_redis.hgetall(name)
        output = {}
        for k, v in r.items():
            output[k.decode('utf8')] = pickle.loads(v)
        return output

    def all(self):
        keys = self.config_redis.keys('*')
        if keys:
            result = {}
            for key in keys:
                if key != f'{self.name}:____':
                    v = self.get_backend(key)
                    result[key] = v
            return result
        else:
            return {}

    # def get(self, name: str):
    #     return self.all().get(name)