from __future__ import annotations

import asyncio
import time
from uuid import uuid4

from redis.asyncio import Redis

from app.services.chat_queue_service import ChatQueueConfig, QueueRequest
from app.services.redis_client import get_redis_client


class RedisChatQueueManager:
    _active_key = "ai_kb:chat_queue:active"
    _waiting_key = "ai_kb:chat_queue:waiting"
    _request_script = """
local active_key = KEYS[1]
local waiting_key = KEYS[2]
local now = tonumber(ARGV[1])
local token = ARGV[2]
local max_active = tonumber(ARGV[3])
local max_waiting = tonumber(ARGV[4])
local active_expire_at = tonumber(ARGV[5])
local key_ttl = tonumber(ARGV[6])

redis.call('ZREMRANGEBYSCORE', active_key, '-inf', now)
local active_count = redis.call('ZCARD', active_key)
local waiting_count = redis.call('LLEN', waiting_key)

if active_count < max_active and waiting_count == 0 then
  redis.call('ZADD', active_key, active_expire_at, token)
  redis.call('EXPIRE', active_key, key_ttl)
  return {'acquired', token, 0, active_count + 1, 0}
end

if waiting_count >= max_waiting then
  redis.call('EXPIRE', active_key, key_ttl)
  redis.call('EXPIRE', waiting_key, key_ttl)
  return {'rejected', token, 0, active_count, waiting_count}
end

redis.call('RPUSH', waiting_key, token)
redis.call('EXPIRE', waiting_key, key_ttl)
redis.call('EXPIRE', active_key, key_ttl)
return {'queued', token, waiting_count + 1, active_count, waiting_count + 1}
"""
    _promote_script = """
local active_key = KEYS[1]
local waiting_key = KEYS[2]
local now = tonumber(ARGV[1])
local token = ARGV[2]
local max_active = tonumber(ARGV[3])
local active_expire_at = tonumber(ARGV[4])
local key_ttl = tonumber(ARGV[5])

redis.call('ZREMRANGEBYSCORE', active_key, '-inf', now)
local active_count = redis.call('ZCARD', active_key)
local waiting_count = redis.call('LLEN', waiting_key)
local head = redis.call('LINDEX', waiting_key, 0)

if head == token and active_count < max_active then
  redis.call('LPOP', waiting_key)
  redis.call('ZADD', active_key, active_expire_at, token)
  redis.call('EXPIRE', active_key, key_ttl)
  redis.call('EXPIRE', waiting_key, key_ttl)
  return {'acquired', token, 0, active_count + 1, waiting_count - 1}
end

local position = redis.call('LPOS', waiting_key, token)
if not position then
  return {'canceled', token, 0, active_count, waiting_count}
end

redis.call('EXPIRE', active_key, key_ttl)
redis.call('EXPIRE', waiting_key, key_ttl)
return {'queued', token, position + 1, active_count, waiting_count}
"""

    async def request_slot(self, config: ChatQueueConfig) -> QueueRequest:
        client = _require_client()
        now = time.time()
        token = uuid4().hex
        result = await client.eval(
            self._request_script,
            2,
            self._active_key,
            self._waiting_key,
            now,
            token,
            config.max_active_requests,
            config.max_queue_size,
            now + config.active_lease_seconds,
            _key_ttl(config),
        )
        return _queue_request(result)

    async def wait_for_slot(self, token: str, config: ChatQueueConfig) -> QueueRequest:
        client = _require_client()
        deadline = time.monotonic() + config.queue_timeout_seconds
        while True:
            now = time.time()
            result = await client.eval(
                self._promote_script,
                2,
                self._active_key,
                self._waiting_key,
                now,
                token,
                config.max_active_requests,
                now + config.active_lease_seconds,
                _key_ttl(config),
            )
            request = _queue_request(result)
            if request.status in {"acquired", "canceled"}:
                return request

            remaining = deadline - time.monotonic()
            if remaining <= 0:
                await self.cancel_waiter(token)
                return QueueRequest(status="timeout", token=token)
            await asyncio.sleep(min(0.25, remaining))

    async def cancel_waiter(self, token: str) -> None:
        client = _require_client()
        await client.lrem(self._waiting_key, 0, token)

    async def release_slot(self, token: str) -> None:
        client = _require_client()
        await client.zrem(self._active_key, token)


redis_chat_queue_manager = RedisChatQueueManager()


def _require_client() -> Redis:
    client = get_redis_client()
    if client is None:
        raise RuntimeError("Redis is not configured")
    return client


def _queue_request(result: list) -> QueueRequest:
    status = str(result[0])
    token = str(result[1]) if len(result) > 1 and result[1] else None
    return QueueRequest(
        status=status,
        token=token,
        position=int(result[2]) if len(result) > 2 else 0,
        active_count=int(result[3]) if len(result) > 3 else 0,
        waiting_count=int(result[4]) if len(result) > 4 else 0,
    )


def _key_ttl(config: ChatQueueConfig) -> int:
    return max(config.active_lease_seconds, config.queue_timeout_seconds) + 300
