"""Redis lease primitives for account ownership (refactor §2).

A sync worker must hold the lease ``sync:lease:{account_id}`` before it runs an
account's sync loops, and no two workers may hold the same lease. Claims are
atomic (``SET NX EX``); renewal and release are owner-checked Lua so a worker can
never renew or delete a lease that has been taken over by another worker.

The set ``sync:live`` is the registry of accounts that *should* be synced; the
claimer only ever claims ids drawn from it.
"""

from __future__ import annotations

import uuid

import redis.asyncio as aioredis

from app.redis_client import aw

LIVE_SET = "sync:live"
_LEASE_PREFIX = "sync:lease:"
WORKERS_PREFIX = "sync:workers:"
_TIER_PREFIX = "sync:tier:"
_SESSION_PREFIX = "sess:active:"


def lease_key(account_id: uuid.UUID) -> str:
    return f"{_LEASE_PREFIX}{account_id}"


def worker_key(worker_id: str) -> str:
    return f"{WORKERS_PREFIX}{worker_id}"


def tier_key(account_id: uuid.UUID) -> str:
    return f"{_TIER_PREFIX}{account_id}"


def session_key(user_id: uuid.UUID | str) -> str:
    """Presence key the API refreshes while a user has an active session (§5)."""
    return f"{_SESSION_PREFIX}{user_id}"


async def mark_user_active(redis: aioredis.Redis, user_id: uuid.UUID | str, grace_s: int) -> None:
    """Refresh the user's session-heat key; workers read it for tier promotion."""
    await aw(redis.set(session_key(user_id), "1", ex=grace_s))


# Renew only if we still own the lease (guards against a stale worker extending a
# lease another worker has already claimed after expiry). Returns 1 on success.
_RENEW_LUA = """
if redis.call('GET', KEYS[1]) == ARGV[1] then
  return redis.call('EXPIRE', KEYS[1], ARGV[2])
else
  return 0
end
"""

# Delete only if we still own the lease. Returns 1 if we released it, else 0.
_RELEASE_LUA = """
if redis.call('GET', KEYS[1]) == ARGV[1] then
  return redis.call('DEL', KEYS[1])
else
  return 0
end
"""


class LeaseScripts:
    """Server-registered Lua for owner-checked renew/release."""

    def __init__(self, redis: aioredis.Redis) -> None:
        self._renew = redis.register_script(_RENEW_LUA)
        self._release = redis.register_script(_RELEASE_LUA)

    async def renew(self, account_id: uuid.UUID, worker_id: str, ttl: int) -> bool:
        result = await self._renew(keys=[lease_key(account_id)], args=[worker_id, ttl])
        return bool(result)

    async def release(self, account_id: uuid.UUID, worker_id: str) -> bool:
        result = await self._release(keys=[lease_key(account_id)], args=[worker_id])
        return bool(result)


async def claim(redis: aioredis.Redis, account_id: uuid.UUID, worker_id: str, ttl: int) -> bool:
    """Atomically claim the lease for ``account_id``; True iff we now own it."""
    ok = await redis.set(lease_key(account_id), worker_id, nx=True, ex=ttl)
    return bool(ok)


async def mark_live(redis: aioredis.Redis, account_id: uuid.UUID) -> None:
    """Register an account as one that should be synced (idempotent)."""
    await aw(redis.sadd(LIVE_SET, str(account_id)))


async def mark_not_live(redis: aioredis.Redis, account_id: uuid.UUID) -> None:
    """Deregister an account and drop any lease so its owner stops it (fencing)."""
    await aw(redis.srem(LIVE_SET, str(account_id)))
    await aw(redis.delete(lease_key(account_id)))