# -*- coding: utf-8 -*-
"""
@File    : rate_limited_pool.py
@Author  : Jay
@Date    : 2026/3/18 17:46
@Desc    : Data collection - general identity token rate limiting
"""
import asyncio
import json
import random

from redis.commands.search.field import TextField, NumericField
from redis.commands.search.index_definition import IndexDefinition, IndexType

from .awaitable import awaitable
from .rtypes import LimitedModel


class RedisearchNotAvailableError(Exception):
    """Raised when RediSearch module is not installed or Redis version is too low."""
    pass


class Limited(object):
    """Rate-limited identity container manager using Redis + RediSearch.

    This class provides synchronous and asynchronous methods for managing
    identity containers with rate limiting enforced via Redis and RediSearch.
    """

    _redisearch_client_db0 = None
    _redis_key = "crawler:identity:flow_limit"

    _lua_ask_script = """
local t = redis.call("TIME")
local now = t[1] * 1000 + math.floor(t[2] / 1000)

local index = ARGV[1]        -- Redisearch index name
local quantity = tonumber(ARGV[2])
local result = {}

local query = string.format("@next_time_available:[-inf %d] | @interval_ms:[0 0]", now)
local search_res = redis.call(
    "FT.SEARCH", index, query, "NOCONTENT",
    "SORTBY", "usage_count", "ASC",
    "LIMIT", 0, quantity
)
if #search_res <= 1 then return nil end

for i = 2, #search_res do
    local key = search_res[i]
    local fields = redis.call("HMGET", key, "id", "containers", "interval_ms")
    local id_val = fields[1]
    local containers = fields[2]
    local interval_ms = tonumber(fields[3] or 0)

    -- Atomically increment usage_count
    local new_usage_count = redis.call("HINCRBY", key, "usage_count", 1)

    -- Update next_time_available if interval_ms > 0 in a single HSET call
    if interval_ms > 0 then
        redis.call("HSET", key, "next_time_available", now + interval_ms)
    end

    -- Insert as a Lua table (will be JSON encoded at the end)
    table.insert(result, {
        id = id_val,
        containers = containers,
        usage_count = new_usage_count
    })
end

return cjson.encode(result)
    """
    _lua_set_script = """
local prefix = ARGV[1]
local json_list = cjson.decode(ARGV[2])

for i, data in ipairs(json_list) do
    local key = prefix .. data["id"]
    local ttl_ms = tonumber(data["ttl_ms"] or 0)

    -- Prepare all fields except ttl_ms for a single HSET call
    local hm = {}
    for field, value in pairs(data) do
        if field ~= "ttl_ms" then
            table.insert(hm, field)
            table.insert(hm, tostring(value))
        end
    end

    -- Write all fields in one HSET
    if #hm > 0 then
        redis.call("HSET", key, unpack(hm))
    end

    -- Set expiration time if ttl_ms > 0
    if ttl_ms > 0 then
        redis.call("PEXPIRE", key, ttl_ms)
    end
end

return 1
    """
    _lua_random_container = """
local index = ARGV[1]
local quantity = tonumber(ARGV[2]) or 1

-- Perform RediSearch query to get all non-empty containers
local search_res = redis.call(
    "FT.SEARCH", index, "*",
    "NOCONTENT",
    "SORTBY", "usage_count", "ASC",
    "LIMIT", 0, quantity
)
if #search_res <= 1 then return nil end

-- Build candidate list (skip first element which is total count)
local candidates = {}
for i = 2, #search_res do
    table.insert(candidates, search_res[i])
end

local result = {}
for i = 2, #search_res do
    local key = search_res[i]

    local fields = redis.call("HMGET", key, "id", "containers", "interval_ms")
    local id_val = fields[1]
    local containers = fields[2]
    local interval_ms = tonumber(fields[3] or 0)

    local new_usage_count = redis.call("HINCRBY", key, "usage_count", 1)
    table.insert(result, {
        id = id_val,
        containers = containers,
        interval_ms = interval_ms,
        usage_count = new_usage_count
    })
end

-- Return JSON-encoded list of containers
return cjson.encode(result)
"""

    def __init__(self, redis_xsync, redis_key=None):
        """Initialize the Limited container manager.

        Args:
            redis_xsync: Redis client factory supporting sync/async context.
            redis_key (str, optional): Redis key prefix for rate-limited containers.
        """
        self._redis_xsync = redis_xsync
        self._redis_key = redis_key or self._redis_key

        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)

        # Verify Redis supports RediSearch module
        self._is_redisearch_supported()

    def _is_redisearch_supported(self) -> bool:
        """Check if RediSearch module is available and create index if needed.

        Raises:
            RedisearchNotAvailableError: If RediSearch is missing or unsupported.

        Returns:
            bool: True if RediSearch is supported.
        """
        with self._redis_xsync(db=0) as redis:
            modules = redis.execute_command("MODULE", "LIST")
            for module in modules:
                for name in module:
                    if isinstance(name, bytes):
                        name = name.decode()
                    if str(name).endswith("redisearch.so"):
                        self._create_index(redis=redis)  # Create index
                        return True
            raise RedisearchNotAvailableError(
                "RediSearch is not installed or your Redis version is too low to support it. "
                "Starting with Redis 8, RediSearch (Redis Query Engine) is integrated and does not need a separate module; "
                "for older versions you must load the module manually. See https://redis.io/docs/latest/develop/ai/search-and-query/ for details."
            )

    def _create_index(self, redis):
        """Create RediSearch index for container rate limiting.

        Args:
            redis: Redis connection object.
        """
        ft = redis.ft(self._redis_key)
        schema = [
            TextField(name="id", sortable=True),
            NumericField(name="ct", sortable=True),
            NumericField(name="ttl_ms", sortable=True),
            NumericField(name="next_time_available", sortable=True),
            NumericField(name="interval_ms", sortable=True),
            NumericField(name="usage_count", sortable=True),
            TextField(name="containers")
        ]
        definition = IndexDefinition(
            prefix=[f"{self._redis_key}:"],
            index_type=IndexType.HASH
        )

        with self._redis_xsync.AsyncRedisAtomicMultiLock(
                "_temporary_index_redis_search",
                ttl_ms=60 * 1000,
                blocking=True,
                blocking_timeout=60 * 1000,
                sleep=0.3 + random.random() * 0.1,
        ):
            existing_indexes = [
                idx.decode() if isinstance(idx, bytes) else idx
                for idx in redis.execute_command("FT._LIST")
            ]
            if self._redis_key in existing_indexes:
                return

            # redis.execute_command("FT.DROPINDEX", self._redis_key)
            ft.create_index(schema, definition=definition)

    # <editor-fold desc="Get N available containers">
    def _ask_available_containers_sync(self, quantity):
        """Synchronously retrieve N available containers from Redis.

        Args:
            quantity (int): Number of containers to retrieve.

        Returns:
            str or None: JSON-encoded container data or None if none available.
        """
        with self._redis_xsync(db=0) as redis:
            encoded_result = redis.eval(self._lua_ask_script, 0, self._redis_key, quantity)
            return encoded_result

    @awaitable(_ask_available_containers_sync)
    async def ask_available_containers(self, quantity):
        """Asynchronously retrieve N available containers from Redis.

        Args:
            quantity (int): Number of containers to retrieve.

        Returns:
            str or None: JSON-encoded container data or None if none available.
        """
        async with self._redis_xsync(db=0) as redis:
            encoded_result = await redis.eval(self._lua_ask_script, 0, self._redis_key, quantity)
            return encoded_result

    def _ask_available_container_sync(self):
        """Synchronously get a single available container.

        Returns:
            str or None: JSON-encoded single container data or None.
        """
        return self._ask_available_containers_sync(quantity=1)

    @awaitable(_ask_available_container_sync)
    async def ask_available_container(self):
        """Asynchronously get a single available container.

        Returns:
            str or None: JSON-encoded single container data or None.
        """
        return await self.ask_available_containers(quantity=1)

    # </editor-fold>

    # <editor-fold desc="Wait until containers are available">
    async def _wait_ask_available_containers(self, quantity=1):
        """Async loop that waits until N containers are available.

        Args:
            quantity (int): Number of containers to wait for.

        Returns:
            str: JSON-encoded container data.
        """
        while True:
            containers = await self.ask_available_containers(quantity=quantity)
            if containers:
                return containers
            await asyncio.sleep(0.3 + random.random() * 0.1)

    def _wait_ask_available_containers_sync(self, quantity=1, timeout=60):
        """Synchronously wait for N containers with timeout.

        Args:
            quantity (int): Number of containers to wait for.
            timeout (int): Timeout in seconds.

        Returns:
            str: JSON-encoded container data.
        """
        return self._loop.run_until_complete(
            asyncio.wait_for(
                self._wait_ask_available_containers(quantity),
                timeout=timeout
            )
        )

    @awaitable(_wait_ask_available_containers_sync)
    async def wait_ask_available_containers(self, quantity=1, timeout=60):
        """Async wait for N containers to become available.

        Args:
            quantity (int): Number of members to retrieve.
            timeout (int): Maximum waiting time in seconds.

        Returns:
            str or None: JSON-encoded container data or None if timeout.
        """
        return await asyncio.wait_for(self._wait_ask_available_containers(quantity), timeout=timeout)

    async def _wait_ask_available_container(self):
        """Async loop that waits until N containers are available.

        Args:
            quantity (int): Number of containers to wait for.

        Returns:
            str: JSON-encoded container data.
        """
        while True:
            containers = await self.ask_available_containers(quantity=1)
            if containers:
                return containers
            await asyncio.sleep(0.3 + random.random() * 0.1)

    def _wait_ask_available_container_sync(self, timeout=60):
        """Synchronously wait for N containers with timeout.

        Args:
            quantity (int): Number of containers to wait for.
            timeout (int): Timeout in seconds.

        Returns:
            str: JSON-encoded container data.
        """
        return self._loop.run_until_complete(
            asyncio.wait_for(
                self._wait_ask_available_container(),
                timeout=timeout
            )
        )

    @awaitable(_wait_ask_available_container_sync)
    async def wait_ask_available_container(self, timeout=60):
        """Async wait for N containers to become available.

        Args:
            quantity (int): Number of members to retrieve.
            timeout (int): Maximum waiting time in seconds.

        Returns:
            str or None: JSON-encoded container data or None if timeout.
        """
        return await asyncio.wait_for(self._wait_ask_available_container(), timeout=timeout)

    # </editor-fold>

    # <editor-fold desc="Set available containers">
    def _set_available_containers_sync(self, *containers: LimitedModel):
        """Synchronously store container configuration in Redis.

        Args:
            containers (LimitedModel): Variable number of container configurations.

        Returns:
            bool: True if operation succeeded.
        """
        with self._redis_xsync(db=0) as redis:
            new_containers = json.dumps([token.model_dump() for token in containers])
            return bool(redis.eval(self._lua_set_script, 0, f"{self._redis_key}:", new_containers))

    @awaitable(_set_available_containers_sync)
    async def set_available_containers(self, *containers: LimitedModel):
        """Asynchronously store container configuration in Redis.

        Args:
            containers (LimitedModel): Variable number of container configurations.

        Returns:
            bool: True if operation succeeded.
        """
        async with self._redis_xsync(db=0) as redis:
            new_containers = json.dumps([token.model_dump() for token in containers])
            return bool(await redis.eval(self._lua_set_script, 0, f"{self._redis_key}:", new_containers))

    # </editor-fold>

    # <editor-fold desc="Randomly acquire a container, ignoring rate limit">
    def _acquire_random_containers_sync(self, quantity=1):
        """Synchronously acquire a specified number of random containers.

        Args:
            quantity (int): Number of containers to retrieve.

        Returns:
            str: JSON-encoded list of container data.
        """
        with self._redis_xsync(db=0) as redis:
            return redis.eval(self._lua_random_container, 0, self._redis_key, quantity)

    # Asynchronous version: acquire N random containers
    @awaitable(_acquire_random_containers_sync)
    async def acquire_random_containers(self, quantity=1):
        """Asynchronously acquire a specified number of random containers.

        Args:
            quantity (int): Number of containers to retrieve.

        Returns:
            str: JSON-encoded list of container data.
        """
        async with self._redis_xsync(db=0) as redis:
            return await redis.eval(self._lua_random_container, 0, self._redis_key, quantity)

    # Synchronous helper to acquire a single random container
    def _acquire_random_container_sync(self):
        """Synchronously acquire a single random container.

        Returns:
            str: JSON-encoded single container data.
        """
        return self.acquire_random_containers(quantity=1)

    # Asynchronous helper to acquire a single random container
    @awaitable(_acquire_random_container_sync)
    async def acquire_random_container(self):
        """Asynchronously acquire a single random container.

        Returns:
            str: JSON-encoded single container data.
        """
        return await self.acquire_random_containers(quantity=1)

    async def _wait_acquire_random_container(self):
        """Async loop that waits until a random container is available.

        Returns:
            str: JSON-encoded single container data.
        """
        while True:
            containers = await self.acquire_random_container()
            if containers: return containers
            await asyncio.sleep(0.3 + random.random() * 0.1)

    def _wait_acquire_random_container_sync(self, timeout=60):
        """Synchronously wait for a random container with timeout.

        Args:
            timeout (int): Timeout in seconds.

        Returns:
            str: JSON-encoded single container data.
        """
        return self._loop.run_until_complete(
            asyncio.wait_for(
                self._wait_acquire_random_container(),
                timeout=timeout
            )
        )

    @awaitable(_wait_acquire_random_container_sync)
    async def wait_acquire_random_container(self, timeout=60):
        """Asynchronously wait for a random container to become available.

        Args:
            timeout (int): Timeout in seconds.

        Returns:
            str: JSON-encoded single container data.
        """
        return await asyncio.wait_for(self._wait_acquire_random_container(), timeout=timeout)

    async def _wait_acquire_random_containers(self, quantity):
        """Async loop that waits until N random containers are available.

        Args:
            quantity (int): Number of containers to wait for.

        Returns:
            str: JSON-encoded container data.
        """
        while True:
            containers = await self.acquire_random_containers(quantity=quantity)
            if containers: return containers
            await asyncio.sleep(0.3 + random.random() * 0.1)

    def _wait_acquire_random_containers_sync(self, quantity=1, timeout=60):
        """Synchronously wait for N random containers with timeout.

        Args:
            quantity (int): Number of containers to wait for.
            timeout (int): Timeout in seconds.

        Returns:
            str: JSON-encoded container data.
        """
        return self._loop.run_until_complete(
            asyncio.wait_for(
                self._wait_acquire_random_containers(quantity),
                timeout=timeout
            )
        )

    @awaitable(_wait_acquire_random_containers_sync)
    async def wait_acquire_random_containers(self, quantity=1, timeout=60):
        """Asynchronously wait for N random containers to become available.

        Args:
            quantity (int): Number of containers to retrieve.
            timeout (int): Maximum waiting time in seconds.

        Returns:
            str: JSON-encoded container data.
        """
        return await asyncio.wait_for(self._wait_acquire_random_containers(quantity), timeout=timeout)

    # </editor-fold>

    async def __aenter__(self):
        """Async context enter.

        Returns:
            Limited: self
        """
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context exit."""
        self.close()

    def __enter__(self):
        """Sync context enter.

        Returns:
            Limited: self
        """
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Sync context exit."""
        self.close()

    def close(self):
        """Close internal asyncio loop."""
        self._loop.close()


__all__ = ["Limited"]
