"""Capacity stores. Redis is the distributed authority; memory is the local fallback."""

from __future__ import annotations

import json
from typing import Protocol

from src.core.provider_capacity.algorithm import (
    apply_observation,
    cancel_waiter,
    release,
    try_reserve,
)
from src.core.provider_capacity.types import (
    BucketSnapshot,
    CapacityBucket,
    CapacityObservation,
    CapacityPolicy,
    ReserveDecision,
)

_UNLOCK = """
if redis.call('GET', KEYS[1]) == ARGV[1] then
  return redis.call('DEL', KEYS[1])
end
return 0
"""


class CapacityStore(Protocol):
    async def try_reserve(
        self,
        bucket: CapacityBucket,
        demand: int,
        waiter_id: str,
        now_ms: int,
        policy: CapacityPolicy,
    ) -> tuple[ReserveDecision, BucketSnapshot]: ...

    async def release(self, bucket: CapacityBucket, demand: int) -> BucketSnapshot: ...

    async def cancel_waiter(self, bucket: CapacityBucket, waiter_id: str) -> None: ...

    async def reconcile(
        self, bucket: CapacityBucket, observation: CapacityObservation, now_ms: int
    ) -> BucketSnapshot: ...


class MemoryCapacityStore:
    """Process-local store. Used in tests and as Redis-outage fallback."""

    def __init__(self) -> None:
        import asyncio

        self._states: dict[str, BucketSnapshot] = {}
        self._lock = asyncio.Lock()

    async def try_reserve(
        self,
        bucket: CapacityBucket,
        demand: int,
        waiter_id: str,
        now_ms: int,
        policy: CapacityPolicy,
    ) -> tuple[ReserveDecision, BucketSnapshot]:
        async with self._lock:
            state = self._states.setdefault(bucket.id, BucketSnapshot())
            decision = try_reserve(
                state, demand=demand, waiter_id=waiter_id, now_ms=now_ms, policy=policy
            )
            return decision, _copy(state)

    async def release(self, bucket: CapacityBucket, demand: int) -> BucketSnapshot:
        async with self._lock:
            state = self._states.setdefault(bucket.id, BucketSnapshot())
            release(state, demand)
            return _copy(state)

    async def cancel_waiter(self, bucket: CapacityBucket, waiter_id: str) -> None:
        async with self._lock:
            state = self._states.setdefault(bucket.id, BucketSnapshot())
            cancel_waiter(state, waiter_id)

    async def reconcile(
        self, bucket: CapacityBucket, observation: CapacityObservation, now_ms: int
    ) -> BucketSnapshot:
        async with self._lock:
            state = self._states.setdefault(bucket.id, BucketSnapshot())
            apply_observation(state, observation, now_ms)
            return _copy(state)


class RedisCapacityStore:
    """Atomic reservation across workers via a Redis mutex around the shared snapshot."""

    def __init__(self, redis_url: str, *, redis_factory=None) -> None:
        self._redis_url = redis_url
        self._redis_factory = redis_factory

    async def try_reserve(
        self,
        bucket: CapacityBucket,
        demand: int,
        waiter_id: str,
        now_ms: int,
        policy: CapacityPolicy,
    ) -> tuple[ReserveDecision, BucketSnapshot]:
        async with _RedisClient(self) as redis:
            async with _BucketLock(redis, bucket):
                state = await _load(redis, bucket)
                decision = try_reserve(
                    state, demand=demand, waiter_id=waiter_id, now_ms=now_ms, policy=policy
                )
                await _save(redis, bucket, state)
                return decision, _copy(state)

    async def release(self, bucket: CapacityBucket, demand: int) -> BucketSnapshot:
        async with _RedisClient(self) as redis:
            async with _BucketLock(redis, bucket):
                state = await _load(redis, bucket)
                release(state, demand)
                await _save(redis, bucket, state)
                return _copy(state)

    async def cancel_waiter(self, bucket: CapacityBucket, waiter_id: str) -> None:
        async with _RedisClient(self) as redis:
            async with _BucketLock(redis, bucket):
                state = await _load(redis, bucket)
                cancel_waiter(state, waiter_id)
                await _save(redis, bucket, state)

    async def reconcile(
        self, bucket: CapacityBucket, observation: CapacityObservation, now_ms: int
    ) -> BucketSnapshot:
        async with _RedisClient(self) as redis:
            async with _BucketLock(redis, bucket):
                state = await _load(redis, bucket)
                apply_observation(state, observation, now_ms)
                await _save(redis, bucket, state)
                return _copy(state)


def _copy(state: BucketSnapshot) -> BucketSnapshot:
    return BucketSnapshot.from_dict(state.to_dict())


def _state_key(bucket: CapacityBucket) -> str:
    return f"provider-capacity:state:{bucket.id}"


def _lock_key(bucket: CapacityBucket) -> str:
    return f"provider-capacity:lock:{bucket.id}"


async def _load(redis, bucket: CapacityBucket) -> BucketSnapshot:
    raw = await redis.get(_state_key(bucket))
    if not raw:
        return BucketSnapshot()
    if isinstance(raw, bytes):
        raw = raw.decode()
    return BucketSnapshot.from_dict(json.loads(raw))


async def _save(redis, bucket: CapacityBucket, state: BucketSnapshot) -> None:
    await redis.set(_state_key(bucket), json.dumps(state.to_dict()), ex=300)


class _RedisClient:
    def __init__(self, store: RedisCapacityStore):
        self._store = store
        self._client = None

    async def __aenter__(self):
        if self._store._redis_factory is not None:
            self._client = self._store._redis_factory()
            if hasattr(self._client, "__aenter__"):
                return await self._client.__aenter__()
            return self._client
        from redis.asyncio import Redis

        self._client = Redis.from_url(
            self._store._redis_url, socket_timeout=2, socket_connect_timeout=2
        )
        return await self._client.__aenter__()

    async def __aexit__(self, exc_type, exc, tb):
        if self._client is None:
            return None
        if hasattr(self._client, "__aexit__"):
            return await self._client.__aexit__(exc_type, exc, tb)
        close = getattr(self._client, "aclose", None) or getattr(self._client, "close", None)
        if close is not None:
            result = close()
            if hasattr(result, "__await__"):
                await result
        return None


class _BucketLock:
    def __init__(self, redis, bucket: CapacityBucket):
        self._redis = redis
        self._key = _lock_key(bucket)
        self._token = ""

    async def __aenter__(self):
        import asyncio
        from uuid import uuid4

        self._token = uuid4().hex
        for _ in range(50):
            acquired = await self._redis.set(self._key, self._token, nx=True, px=5000)
            if acquired:
                return self
            await asyncio.sleep(0.02)
        raise TimeoutError("provider_capacity_lock_timeout")

    async def __aexit__(self, exc_type, exc, tb):
        await self._redis.eval(_UNLOCK, 1, self._key, self._token)
        return None