import asyncio
import json

import redis.asyncio as aioredis
from redis.exceptions import ResponseError

# One server-side operation: a successful write always has an expiry.
_SET_JSON = """
local kind = redis.call('TYPE', KEYS[1]).ok
if kind == 'string' or kind == 'list' then
  redis.call('DEL', KEYS[1])
end
redis.call('JSON.SET', KEYS[1], '$', ARGV[1])
redis.call('EXPIRE', KEYS[1], ARGV[2])
return 1
"""

_SET_LIST = """
redis.call('DEL', KEYS[1])
for i = 2, #ARGV do
  redis.call('RPUSH', KEYS[1], ARGV[i])
end
redis.call('EXPIRE', KEYS[1], ARGV[1])
return 1
"""

class AsyncCache:
  """Native strings/string lists and typed JSON, with expiry on every write."""

  def __init__(self, backend="memory", prefix="cache:", default_ttl=300, **kwargs):
    self.prefix = prefix
    self.ttl = default_ttl
    if backend == "redis":
      self.backend_type = "redis"
      self.redis = aioredis.Redis(
        host=kwargs.get("host", "localhost"),
        port=int(kwargs.get("port", 6379)),
        username=kwargs.get("username"),
        password=kwargs.get("password"),
        db=int(kwargs.get("db", 0)),
        decode_responses=False,
      )
    else:
      self.backend_type = "memory"
      self._store = {}
      self._expiry = {}

  def _key(self, key):
    return f"{self.prefix}{key}"

  async def set(self, key, value, ttl=None):
    ttl = self.ttl if ttl is None else ttl
    if isinstance(ttl, bool) or not isinstance(ttl, int) or ttl <= 0:
      raise ValueError("Cache TTL must be a positive integer")
    dumped = json.dumps(value, ensure_ascii=False, allow_nan=False)
    full_key = self._key(key)
    if self.backend_type == "redis":
      if isinstance(value, str):
        # SET replaces either a string or a JSON key and sets TTL atomically.
        return bool(await self.redis.set(full_key, value, ex=ttl))
      if isinstance(value, list) and value and all(isinstance(item, str) for item in value):
        return bool(await self.redis.eval(_SET_LIST, 1, full_key, ttl, *value))
      return bool(await self.redis.eval(_SET_JSON, 1, full_key, dumped, ttl))
    self._store[full_key] = dumped
    self._expiry[full_key] = asyncio.get_running_loop().time() + ttl
    return True

  async def get(self, key):
    full_key = self._key(key)
    if self.backend_type == "redis":
      # JSON.GET also reads existing JSON scalar values without migration.
      try:
        data = await self.redis.execute_command("JSON.GET", full_key)
      except ResponseError as exc:
        if not str(exc).startswith("WRONGTYPE"):
          raise
        try:
          data = await self.redis.get(full_key)
        except ResponseError as string_error:
          if not str(string_error).startswith("WRONGTYPE"):
            raise
          items = await self.redis.lrange(full_key, 0, -1)
          return [item.decode("utf-8") if isinstance(item, bytes) else item for item in items] or None
        if data is None:
          return None
        try:
          # Raw strings must not be JSON-decoded: '123' remains a string.
          return data.decode("utf-8") if isinstance(data, bytes) else data
        except UnicodeDecodeError:
          # Old binary Pickle entries are cache misses, never deserialized.
          return None
    else:
      expiry = self._expiry.get(full_key)
      if expiry is not None and expiry <= asyncio.get_running_loop().time():
        self._store.pop(full_key, None)
        self._expiry.pop(full_key, None)
        return None
      data = self._store.get(full_key)
    return None if data is None else json.loads(data)

  async def delete(self, key):
    full_key = self._key(key)
    if self.backend_type == "redis":
      return await self.redis.delete(full_key)
    self._store.pop(full_key, None)
    self._expiry.pop(full_key, None)
    return True

  async def clear(self):
    if self.backend_type == "redis":
      count = 0
      async for key in self.redis.scan_iter(match=f"{self.prefix}*"):
        count += await self.redis.delete(key)
      return count
    self._store.clear()
    self._expiry.clear()
    return True