"""Redis helpers: API cache (prefix `aia:api:`), distributed locks (`aia:lock:`), heartbeats."""
from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any

import orjson
from redis.asyncio import Redis

from aiatlas.config import settings

log = logging.getLogger(__name__)
_redis: Redis | None = None


def redis() -> Redis:
    global _redis
    if _redis is None:
        _redis = Redis.from_url(settings.redis_url, decode_responses=False, socket_connect_timeout=3, socket_timeout=5)
    return _redis


async def cache_get(key: str) -> Any | None:
    try:
        raw = await redis().get(f"aia:api:{key}")
        return orjson.loads(raw) if raw else None
    except Exception:  # noqa: BLE001
        return None


async def cache_set(key: str, value: Any, ttl_s: int = 300) -> None:
    try:
        await redis().set(f"aia:api:{key}", orjson.dumps(value, default=str), ex=ttl_s)
    except Exception:  # noqa: BLE001
        pass


async def cache_invalidate(prefix: str = "") -> int:
    n = 0
    try:
        async for key in redis().scan_iter(match=f"aia:api:{prefix}*", count=500):
            await redis().delete(key)
            n += 1
    except Exception:  # noqa: BLE001
        pass
    return n


@asynccontextmanager
async def lock(name: str, ttl_s: int = 3600) -> AsyncIterator[bool]:
    """Best-effort distributed lock. Yields False when someone else holds it (callers must skip)."""
    key = f"aia:lock:{name}"
    acquired = False
    try:
        acquired = bool(await redis().set(key, b"1", nx=True, ex=ttl_s))
    except Exception as exc:  # noqa: BLE001
        log.warning("redis unavailable, proceeding without lock", extra={"error": str(exc)})
        acquired = True
    try:
        yield acquired
    finally:
        if acquired:
            try:
                await redis().delete(key)
            except Exception:  # noqa: BLE001
                pass


async def heartbeat(service: str, payload: dict[str, Any], ttl_s: int = 180) -> None:
    try:
        await redis().set(f"aia:heartbeat:{service}", orjson.dumps(payload, default=str), ex=ttl_s)
    except Exception:  # noqa: BLE001
        pass


async def heartbeats() -> dict[str, Any]:
    out: dict[str, Any] = {}
    try:
        async for key in redis().scan_iter(match="aia:heartbeat:*"):
            raw = await redis().get(key)
            if raw:
                out[key.decode().rsplit(":", 1)[-1]] = orjson.loads(raw)
    except Exception:  # noqa: BLE001
        pass
    return out


async def close() -> None:
    global _redis
    if _redis is not None:
        await _redis.aclose()
        _redis = None
    await asyncio.sleep(0)


__all__ = ["cache_get", "cache_invalidate", "cache_set", "close", "heartbeat", "heartbeats", "lock", "redis"]