from __future__ import annotations

import asyncio
import datetime
import os
import struct
import typing
from itertools import chain

from loguru import logger
from redis.asyncio import ConnectionPool as AsyncConnectionPool
from redis.asyncio import Redis as AsyncRedisBase
from redis.backoff import ExponentialBackoff
from redis.exceptions import BusyLoadingError, ConnectionError, TimeoutError
from redis.retry import Retry

from good_common.dependencies import BaseProvider
from redis import ConnectionPool, StrictRedis

_Key = typing.Union[str, bytes, int, float]
_Value = typing.Union[str, bytes, int, float, bool, None]


_script_psadd = """
local _target_key = KEYS[1]
local _name = ARGV[1]
local _value = ARGV[2]
local _score = ARGV[3]

local hash_table = _name .. ':t'
local sorted_set = _name .. ':s'
local exclusion_set = _name .. ':exc'
local expiration_set = _name .. ':exp'
                                
redis.call('SADD', ':ps:keys', _name)
                                
if redis.call('SISMEMBER', exclusion_set, _target_key) == 0 and redis.call('ZSCORE', expiration_set, _target_key) == false then
    redis.call('HSET', hash_table, _target_key, _value)
    return tonumber(redis.call('ZINCRBY', sorted_set, _score, _target_key))
else
    return tonumber(0)
end
"""

_script_pspull = """
local _name = ARGV[1]
local _count = ARGV[2]
local _purge = ARGV[3]
local _gate = ARGV[4]
local _expire_in = ARGV[5] or 1

local sorted_set = _name .. ':s'
local expiration_set = _name .. ':exp'
local hash_table = _name .. ':t'
local exclusion_set = _name .. ':exc'

local output = {}
for i = 1, tonumber(_count) do
    local result = redis.call('ZPOPMAX', sorted_set, 1)
    if next(result) == nil then break end
    local key = result[1]
    local value = redis.call('HGET', hash_table, key)
    if value then
        table.insert(output, key)
        table.insert(output, value)
        if _purge == '1' then
            redis.call('HDEL', hash_table, key)
            if _gate == '1' then
                redis.call('SADD', exclusion_set, key)
            end
        else
            redis.call('ZADD', sorted_set, -1, key)
        end

        redis.call('ZADD', expiration_set, _expire_in, key)
        
    end
end

return output
"""

_script_psdecrementall = """
local _name = ARGV[1]
local _decrement = ARGV[2]

local sorted_set = _name .. ':s'

local keys = redis.call('ZRANGE', sorted_set, 0, -1)

for i, key in ipairs(keys) do
    redis.call('ZINCRBY', sorted_set, -_decrement, key)
end

return 1
"""


class Redis(StrictRedis):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.__psadd = self.register_script(_script_psadd)
        self.__pspull = self.register_script(_script_pspull)
        self.__psdecrementall = self.register_script(_script_psdecrementall)

    #
    # Prioritized Stream
    #
    def psadd(
        self, name: str, key: _Key, value: _Value, score: int = 1, readd: bool = False
    ):
        if readd:
            self.srem(f"{name}:exc", key)
            self.zrem(f"{name}:exp", key)

        return self.__psadd(keys=[key], args=[name, value, score])

    def psaddmany(
        self, name: str, items: list[tuple[_Key, _Value, int]], readd: bool = False
    ):
        pipe = self.pipeline()

        if readd:
            pipe.srem(f"{name}:exc", *[key for key, _, _ in items])
            pipe.zrem(f"{name}:exp", *[key for key, _, _ in items])

        for key, value, score in items:
            self.__psadd(keys=[key], args=[name, value, score], client=pipe)

        return pipe.execute()

    def pspull(
        self,
        name: str,
        count: int = 1,
        purge: bool = True,
        gate: bool = False,
        expire_at: datetime.timedelta | datetime.datetime = datetime.timedelta(
            minutes=24 * 60
        ),
    ):
        if isinstance(expire_at, datetime.timedelta):
            expire_at = datetime.datetime.now(tz=datetime.UTC) + expire_at

        r = self.__pspull(
            args=[
                name,
                count,
                "1" if purge else "0",
                "1" if gate else "0",
                int(expire_at.timestamp()),
            ]
        )

        _result = {r[i]: r[i + 1] for i in range(0, len(r), 2)}

        return {
            key.decode("utf-8") if isinstance(key, bytes) else key: _result[key]
            for key in _result
        }

    def pskeys(self, name: str, **kwargs):
        return self.zrange(f"{name}:s", 0, -1, **kwargs)

    def psvalues(self, name: str):
        v = self.hgetall(f"{name}:t")
        return {
            key.decode("utf-8") if isinstance(key, bytes) else key: value
            for key, value in v.items()
        }

    def psdelete(self, name: str):
        self.delete(f"{name}:t", f"{name}:s", f"{name}:exc")
        self.srem(":ps:keys", name)

    def psincrement(self, name: str, key: _Key, score: int):
        return self.zincrby(f"{name}:s", score, key)

    def psdeletebelowthreshold(self, name: str, threshold: int, gate: bool = False):
        keys = self.zrangebyscore(f"{name}:s", "-inf", threshold)
        pipe = self.pipeline()
        for key in keys:
            pipe.hdel(f"{name}:t", key)
            pipe.zrem(f"{name}:s", key)
            if gate:
                pipe.sadd(f"{name}:exc", key)
        pipe.execute()
        return keys

    def pslen(self, name: str):
        return self.zcard(f"{name}:s")

    def psget(self, name: str, key: _Key):
        return self.hget(f"{name}:t", key)

    def psdecrement_all(self, name: str, decrement: int):
        return self.__psdecrementall(args=[name, decrement])

    def psexpire(self):
        ts = int(datetime.datetime.now().timestamp())
        # logger.info(ts)
        for key in self.smembers(":ps:keys"):
            key = key.decode("utf-8")
            # logger.info((f'{key}:exp', '-inf', ts))
            expired_keys = self.zrangebyscore(f"{key}:exp", "-inf", ts)
            # logger.info(expired_keys)
            for ekey in expired_keys:
                self.zrem(f"{key}:exp", ekey.decode("utf-8"))
                self.hdel(f"{key}:t", ekey.decode("utf-8"))
                self.zrem(f"{key}:s", ekey.decode("utf-8"))
                # self.sadd(f'{key}:exc', ekey)
            # logger.info(key)
            # for key in self.zrange(f'{key}:s', 0, -1, withscores=True):
            #     if key[1] > ts:
            #         logger.info(key)
            #         self.hdel(f'{key}:t', key[0])
            #         self.zrem(f'{key}:s', key[0])
            #         self.sadd(f'{key}:exc', key[0])

    def _hashany(self, value: typing.Any):
        return struct.pack(">Q", hash(value) & 0xFFFFFFFFFFFFFFFF)

    def hashset_add(self, collection, value: typing.Any):
        self.sadd(collection, self._hashany(value))

    def hashset_remove(self, collection, value: typing.Any):
        self.srem(collection, self._hashany(value))

    def hashset_contains(self, collection, value: typing.Any):
        return self.sismember(collection, self._hashany(value)) == 1

    def hashset_clear(self, collection):
        self.delete(collection)


class RedisProvider(BaseProvider[Redis], Redis):
    instance_cache: dict[tuple[typing.Any, ...], Redis] = {}

    def initializer(self, cls_args, cls_kwargs, fn_kwargs):
        return cls_args, cls_kwargs

    @classmethod
    def provide(cls, *args, **kwargs):
        _key = tuple(chain.from_iterable([args, kwargs.items()]))
        # logger.info(f'Key: {_key}')
        if _key not in cls.instance_cache:
            # logger.info("Creating new Redis instance")
            retry = Retry(ExponentialBackoff(), 3)
            connection_pool = ConnectionPool(
                host=kwargs.get("host", os.environ.get("REDIS_HOST", "localhost")),
                port=kwargs.get("port", int(os.environ.get("REDIS_PORT", "6379"))),
                db=kwargs.get("db", int(os.environ.get("REDIS_DB", "0"))),
                decode_responses=kwargs.get("decode_responses", False),
                protocol=3,
                retry=retry,
                retry_on_error=[BusyLoadingError, ConnectionError, TimeoutError],
            )
            cls.instance_cache[_key] = super().provide(connection_pool=connection_pool)
        return cls.instance_cache[_key]


class AsyncRedis(AsyncRedisBase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.__psadd = self.register_script(_script_psadd)
        self.__pspull = self.register_script(_script_pspull)
        self.__psdecrementall = self.register_script(_script_psdecrementall)

    async def psadd(
        self, name: str, key: _Key, value: _Value, score: int = 1, readd: bool = False
    ):
        async with self.pipeline() as pipe:
            if readd:
                await pipe.srem(f"{name}:exc", key)
                await pipe.zrem(f"{name}:exp", key)

            await self.__psadd(keys=[key], args=[name, value, score], client=pipe)
            r = await pipe.execute()
            if readd:
                return r[-1]
            else:
                return r

    async def psaddmany(
        self, name: str, items: list[tuple[_Key, _Value, int]], readd: bool = False
    ):
        async with self.pipeline() as pipe:
            if readd:
                pipe.srem(f"{name}:exc", *[key for key, _, _ in items])
                pipe.zrem(f"{name}:exp", *[key for key, _, _ in items])

            for key, value, score in items:
                await self.__psadd(keys=[key], args=[name, value, score], client=pipe)

            r = await pipe.execute()
            if readd:
                return r[2:]
            else:
                return r

    async def pspull(
        self,
        name: str,
        count: int = 1,
        purge: bool = True,
        gate: bool = False,
        expire_at: datetime.timedelta | datetime.datetime = datetime.timedelta(
            minutes=24 * 60
        ),
    ):
        if isinstance(expire_at, datetime.timedelta):
            expire_at = datetime.datetime.now(tz=datetime.UTC) + expire_at

        r = await self.__pspull(
            args=[
                name,
                count,
                "1" if purge else "0",
                "1" if gate else "0",
                int(expire_at.timestamp()),
            ]
        )

        _result = {r[i]: r[i + 1] for i in range(0, len(r), 2)}

        return {
            key.decode("utf-8") if isinstance(key, bytes) else key: _result[key]
            for key in _result
        }

    async def pskeys(self, name: str, **kwargs):
        return await self.zrange(f"{name}:s", 0, -1, **kwargs)

    async def psvalues(self, name: str):
        v = await self.hgetall(f"{name}:t")
        return {
            key.decode("utf-8") if isinstance(key, bytes) else key: value
            for key, value in v.items()
        }

    async def psdelete(self, name: str):
        async with self.pipeline() as pipe:
            await pipe.delete(f"{name}:t", f"{name}:s", f"{name}:exc")
            await pipe.rem(":ps:keys", name)
            return await pipe.execute()

    async def psincrement(self, name: str, key: _Key, score: int):
        return await self.zincrby(f"{name}:s", score, key)

    async def psdeletebelowthreshold(
        self, name: str, threshold: int, gate: bool = False
    ):
        async with self.pipeline() as pipe:
            keys = await self.zrangebyscore(f"{name}:s", "-inf", threshold)
            for key in keys:
                await pipe.hdel(f"{name}:t", key)
                await pipe.zrem(f"{name}:s", key)
                if gate:
                    await pipe.sadd(f"{name}:exc", key)
            await pipe.execute()
            return keys

    async def pslen(self, name: str):
        return await self.zcard(f"{name}:s")

    async def psget(self, name: str, key: _Key):
        return await self.hget(f"{name}:t", key)

    async def psdecrement_all(self, name: str, decrement: int):
        return await self.__psdecrementall(args=[name, decrement])

    async def psexpire(self):
        ts = int(datetime.datetime.now().timestamp())
        async with self.pipeline() as pipe:
            # logger.info(ts)
            for key in await self.smembers(":ps:keys"):
                key = key.decode("utf-8")
                # logger.info((f'{key}:exp', '-inf', ts))
                expired_keys = await self.zrangebyscore(f"{key}:exp", "-inf", ts)
                # logger.info(expired_keys)
                for ekey in expired_keys:
                    await self.zrem(f"{key}:exp", ekey.decode("utf-8"))
                    await self.hdel(f"{key}:t", ekey.decode("utf-8"))
                    await self.zrem(f"{key}:s", ekey.decode("utf-8"))

            return await pipe.execute()

    def _hashany(self, value: typing.Any):
        return struct.pack(">Q", hash(value) & 0xFFFFFFFFFFFFFFFF)

    async def hashset_add(self, collection, *values: typing.Any):
        if len(values) == 0:
            return
        _values = [self._hashany(value) for value in values]
        return await self.sadd(collection, *_values)

    async def hashset_remove(self, collection, value: typing.Any):
        return await self.srem(collection, self._hashany(value))

    async def hashset_contains(self, collection, value: typing.Any):
        return (await self.sismember(collection, self._hashany(value))) == 1

    async def hashset_clear(self, collection):
        return await self.delete(collection)


class AsyncRedisClient:
    def __init__(self, pool: AsyncConnectionPool):
        self.pool = pool

    async def __aenter__(self) -> AsyncRedis:
        self.redis = AsyncRedis(connection_pool=self.pool)
        return self.redis

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.redis.close()


class AsyncRedisProvider(BaseProvider[AsyncRedis], AsyncRedis):
    instance_cache: dict[tuple[typing.Any, ...], AsyncRedis] = {}

    def initializer(self, cls_args, cls_kwargs, fn_kwargs):
        return cls_args, cls_kwargs

    @classmethod
    def provide(cls, *args, **kwargs):
        _key = tuple(chain.from_iterable([args, kwargs.items()]))
        # logger.info(f'Key: {_key}')
        if _key not in cls.instance_cache:
            logger.info("Creating new Redis instance")
            retry = Retry(ExponentialBackoff(), 3)
            connection_pool = AsyncConnectionPool(
                host=kwargs.get("host", os.environ.get("REDIS_HOST", "localhost")),
                port=kwargs.get("port", int(os.environ.get("REDIS_PORT", "6379"))),
                db=kwargs.get("db", int(os.environ.get("REDIS_DB", "0"))),
                decode_responses=kwargs.get("decode_responses", False),
                protocol=3,
                retry=retry,
                retry_on_error=[BusyLoadingError, ConnectionError, TimeoutError],
            )
            cls.instance_cache[_key] = super().provide(connection_pool=connection_pool)
        return cls.instance_cache[_key]


class AsyncRedisClientProvider(BaseProvider[AsyncRedisClient], AsyncRedisClient):
    instance_cache: dict[tuple[typing.Any, ...], AsyncConnectionPool] = {}

    def initializer(self, cls_args, cls_kwargs, fn_kwargs):
        return cls_args, cls_kwargs

    @classmethod
    def provide(cls, *args, **kwargs):
        _key = tuple(chain.from_iterable([args, kwargs.items()]))
        # logger.info(f'Key: {_key}')
        if _key not in cls.instance_cache:
            logger.info("Creating new Redis instance")
            retry = Retry(ExponentialBackoff(), 3)
            connection_pool = AsyncConnectionPool(
                host=kwargs.get("host", os.environ.get("REDIS_HOST", "localhost")),
                port=kwargs.get("port", int(os.environ.get("REDIS_PORT", "6379"))),
                db=kwargs.get("db", int(os.environ.get("REDIS_DB", "0"))),
                decode_responses=kwargs.get("decode_responses", False),
                protocol=3,
                retry=retry,
                retry_on_error=[BusyLoadingError, ConnectionError, TimeoutError],
            )
            cls.instance_cache[_key] = connection_pool

        return AsyncRedisClient(pool=cls.instance_cache[_key])

    @classmethod
    async def close(cls):
        await asyncio.gather(*[pool.aclose() for pool in cls.instance_cache.values()])
        cls.instance_cache.clear()
        return True
