"""Native JSON storage for the asyncio Redis session backend."""

import json

async def check_support(backend):
  """Fail at startup rather than silently falling back to string storage."""
  commands = await backend.execute_command("COMMAND", "INFO", "JSON.GET", "JSON.SET")
  entries = list(commands.values()) if isinstance(commands, dict) else commands
  if len(entries) != 2 or not all(entries):
    raise RuntimeError("SESSION_REDIS_JSON requires JSON.GET and JSON.SET support")

async def get(backend, key):
  """Read legacy strings as well as native JSON documents."""
  from redis.exceptions import ResponseError

  try:
    return await backend.get(key)
  except ResponseError as error:
    if not str(error).startswith("WRONGTYPE"):
      raise
    # A single root path returns the object, not a JSONPath result array.
    return await backend.execute_command("JSON.GET", key, ".")

async def set(backend, key, value, expiry):
  """Replace a string or JSON key and apply expiry in one transaction."""
  # Validate before deleting an existing key. Preserve TaggedJSONSerializer's
  # representation so tuples, datetimes, etc. still round-trip normally.
  document = json.loads(value)
  if not isinstance(document, dict):
    raise ValueError("Native Redis JSON sessions require a JSON object")
  async with backend.pipeline(transaction=True) as pipeline:
    pipeline.delete(key)
    pipeline.execute_command("JSON.SET", key, "$", value)
    if expiry is not None:
      pipeline.expire(key, expiry)
    results = await pipeline.execute()
  return results[1]