import logging
from typing import Optional

import rb
import redis_lock

from src.config import Configuration

cluster: Optional[rb.Cluster] = None


def get_cluster() -> rb.Cluster:
    if not cluster:
        raise Exception("Redis cluster not set")
    return cluster


class KemonoRouter(rb.BaseRouter):
    def get_host_for_key(self, key):
        top_level_prefix_of_key = key.split(":")[0]
        if self.key_spaces().get(top_level_prefix_of_key) is not None:
            return self.key_spaces()[top_level_prefix_of_key]
        else:
            logging.error(f"Missing redis key space for {key}, {top_level_prefix_of_key}")
            if Configuration().development_mode:
                return 0
            raise rb.UnroutableCommand()

    @staticmethod
    def key_spaces():
        return Configuration().redis["keyspaces"]

    def get_key(self, command, args):
        from rb._rediscommands import COMMANDS
        from rb.router import extract_keys

        spec = COMMANDS.get(command.upper())

        keys = extract_keys(args, spec["key_spec"])

        if len(keys) > 1:
            args = [args[0]]
        return super().get_key(command, args)


class KemonoRedisLock(redis_lock.Lock):
    """Reword to make the module compatible with Redis-Blaster."""

    def release(self):
        if self._lock_renewal_thread is not None:
            self._stop_lock_renewer()
        # soft reimplementation of UNLOCK_SCRIPT in Python
        self._client.delete(self._signal)
        self._client.lpush(self._signal, 1)
        self._client.pexpire(self._signal, self._signal_expire)
        self._client.delete(self._name)

    def extend(self, expire=None):
        if expire:
            expire = int(expire)
            if expire < 0:
                raise ValueError("A negative expire is not acceptable.")
        elif self._expire is not None:
            expire = self._expire
        else:
            raise TypeError(
                "To extend a lock 'expire' must be provided as an "
                "argument to extend() method or at initialization time."
            )
        # soft reimplementation of EXTEND_SCRIPT in Python
        self._client.expire(self._name, expire)


def init():
    global cluster
    cluster = rb.Cluster(
        host_defaults=Configuration().redis["node_options"],
        hosts=Configuration().redis["nodes"],
        router_cls=KemonoRouter,
    )
    return cluster


def get_conn():
    return get_cluster().get_routing_client()


def scan_keys(pattern, count=50000):
    return cluster.get_local_client_for_key(pattern).scan_iter(match=pattern, count=count)


def set_multiple_expire_keys(keys, ex=Configuration().redis["default_ttl"]):
    """only use if its for the same key space"""

    example_key = keys.keys().__iter__().__next__()
    redis_client = get_cluster().get_local_client_for_key(example_key)
    pipe = redis_client.pipeline()
    for k in keys:
        pipe.expire(k, ex)
    pipe.execute()
    redis_client.close()
