import json

from telebot.storage.base_storage import StateContext, StateStorageBase

redis_installed = True
try:
    from redis import ConnectionPool, Redis

except:
    redis_installed = False


class StateRedisStorage(StateStorageBase):
    """
    This class is for Redis storage.
    This will work only for states.
    To use it, just pass this class to:
    TeleBot(storage=StateRedisStorage())
    """

    def __init__(
        self,
        host="localhost",
        port=6379,
        db=0,
        password=None,
        prefix="telebot_",
        redis_url=None,
    ):
        super().__init__()
        if redis_url:
            self.redis = ConnectionPool.from_url(redis_url)
        else:
            self.redis = ConnectionPool(host=host, port=port, db=db, password=password)
        # self.con = Redis(connection_pool=self.redis) -> use this when necessary
        #
        # {chat_id: {user_id: {'state': None, 'data': {}}, ...}, ...}
        self.prefix = prefix
        if not redis_installed:
            raise Exception(
                "Redis is not installed. Install it via 'pip install redis'"
            )

    def get_record(self, key):
        """
        Function to get record from database.
        It has nothing to do with states.
        Made for backward compatibility
        """
        connection = Redis(connection_pool=self.redis)
        result = connection.get(self.prefix + str(key))
        connection.close()
        if result:
            return json.loads(result)
        return

    def set_record(self, key, value):
        """
        Function to set record to database.
        It has nothing to do with states.
        Made for backward compatibility
        """
        connection = Redis(connection_pool=self.redis)
        connection.set(self.prefix + str(key), json.dumps(value))
        connection.close()
        return True

    def delete_record(self, key):
        """
        Function to delete record from database.
        It has nothing to do with states.
        Made for backward compatibility
        """
        connection = Redis(connection_pool=self.redis)
        connection.delete(self.prefix + str(key))
        connection.close()
        return True

    def set_state(self, chat_id, user_id, state):
        """
        Set state for a particular user in a chat.
        """
        response = self.get_record(chat_id)
        user_id = str(user_id)
        if hasattr(state, "name"):
            state = state.name

        if response:
            if user_id in response:
                response[user_id]["state"] = state
            else:
                response[user_id] = {"state": state, "data": {}}
        else:
            response = {user_id: {"state": state, "data": {}}}
        self.set_record(chat_id, response)

        return True

    def delete_state(self, chat_id, user_id):
        """
        Delete state for a particular user in a chat.
        """
        response = self.get_record(chat_id)
        user_id = str(user_id)
        if response:
            if user_id in response:
                del response[user_id]
                if user_id == str(chat_id):
                    self.delete_record(chat_id)
                    return True
                else:
                    self.set_record(chat_id, response)
                return True
        return False

    def get_value(self, chat_id, user_id, key):
        """
        Get value for a data of a user in a chat.
        """
        response = self.get_record(chat_id)
        user_id = str(user_id)
        if response:
            if user_id in response:
                if key in response[user_id]["data"]:
                    return response[user_id]["data"][key]
        return None

    def get_state(self, chat_id, user_id):
        """
        Get state of a user in a chat.
        """
        response = self.get_record(chat_id)
        user_id = str(user_id)
        if response:
            if user_id in response:
                return response[user_id]["state"]

        return None

    def get_data(self, chat_id, user_id):
        """
        Get data of particular user in a particular chat.
        """
        response = self.get_record(chat_id)
        user_id = str(user_id)
        if response:
            if user_id in response:
                return response[user_id]["data"]
        return None

    def reset_data(self, chat_id, user_id):
        """
        Reset data of a user in a chat.
        """
        response = self.get_record(chat_id)
        user_id = str(user_id)
        if response:
            if user_id in response:
                response[user_id]["data"] = {}
                self.set_record(chat_id, response)
                return True

    def set_data(self, chat_id, user_id, key, value):
        """
        Set data without interactive data.
        """
        response = self.get_record(chat_id)
        user_id = str(user_id)
        if response:
            if user_id in response:
                response[user_id]["data"][key] = value
                self.set_record(chat_id, response)
                return True
        return False

    def get_interactive_data(self, chat_id, user_id):
        """
        Get Data in interactive way.
        You can use with() with this function.
        """
        return StateContext(self, chat_id, user_id)

    def save(self, chat_id, user_id, data):
        response = self.get_record(chat_id)
        user_id = str(user_id)
        if response:
            if user_id in response:
                response[user_id]["data"] = data
                self.set_record(chat_id, response)
                return True
