from __future__ import annotations

import asyncio
import hashlib
import secrets
from datetime import timedelta

from .models import LeaseView, utc_now


class LeaseCapacityExceeded(RuntimeError):
    pass


class IdempotencyConflict(RuntimeError):
    pass


class LeaseStore:
    def __init__(self, *, endpoint: str, max_active: int, max_records: int, ttl_s: float) -> None:
        self.endpoint = endpoint
        self.max_active = max_active
        self.max_records = max_records
        self.ttl_s = ttl_s
        self._leases: dict[str, LeaseView] = {}
        self._idempotency: dict[str, tuple[str, str]] = {}
        self._lock = asyncio.Lock()

    @staticmethod
    def _request_hash(capability: str) -> str:
        return hashlib.sha256(capability.encode("ascii")).hexdigest()

    def _expire_locked(self) -> None:
        now = utc_now()
        for lease in self._leases.values():
            if lease.active and lease.expires_at <= now:
                lease.active = False
                lease.released_at = now

    def _prune_locked(self) -> None:
        if len(self._leases) < self.max_records:
            return
        inactive = sorted(
            (item for item in self._leases.values() if not item.active),
            key=lambda item: item.released_at or item.expires_at,
        )
        while inactive and len(self._leases) >= self.max_records:
            lease = inactive.pop(0)
            self._leases.pop(lease.lease_id, None)
            for key, (_digest, lease_id) in list(self._idempotency.items()):
                if lease_id == lease.lease_id:
                    self._idempotency.pop(key, None)

    async def acquire(self, capability: str, idempotency_key: str) -> tuple[LeaseView, bool]:
        digest = self._request_hash(capability)
        async with self._lock:
            self._expire_locked()
            existing = self._idempotency.get(idempotency_key)
            if existing:
                if existing[0] != digest:
                    raise IdempotencyConflict("idempotency key reused with another capability")
                lease = self._leases.get(existing[1])
                if lease is None:
                    raise IdempotencyConflict("idempotent lease record is no longer retained")
                return lease.model_copy(deep=True), False
            if sum(1 for item in self._leases.values() if item.active) >= self.max_active:
                raise LeaseCapacityExceeded("active lease capacity reached")
            self._prune_locked()
            if len(self._leases) >= self.max_records:
                raise LeaseCapacityExceeded("lease record capacity reached")
            now = utc_now()
            lease = LeaseView(
                lease_id=f"lease_{secrets.token_hex(16)}",
                capability=capability,
                endpoint=self.endpoint,
                created_at=now,
                expires_at=now + timedelta(seconds=self.ttl_s),
            )
            self._leases[lease.lease_id] = lease
            self._idempotency[idempotency_key] = (digest, lease.lease_id)
            return lease.model_copy(deep=True), True

    async def get(self, lease_id: str) -> LeaseView | None:
        async with self._lock:
            self._expire_locked()
            lease = self._leases.get(lease_id)
            return lease.model_copy(deep=True) if lease else None

    async def release(self, lease_id: str) -> tuple[LeaseView | None, bool]:
        async with self._lock:
            self._expire_locked()
            lease = self._leases.get(lease_id)
            if lease is None:
                return None, False
            changed = lease.active
            if changed:
                lease.active = False
                lease.released_at = utc_now()
            return lease.model_copy(deep=True), changed

    async def discard(self, lease_id: str) -> None:
        async with self._lock:
            self._leases.pop(lease_id, None)
            for key, (_digest, stored_id) in list(self._idempotency.items()):
                if stored_id == lease_id:
                    self._idempotency.pop(key, None)

    async def active_count(self, capability: str | None = None) -> int:
        async with self._lock:
            self._expire_locked()
            return sum(
                1
                for item in self._leases.values()
                if item.active and (capability is None or item.capability == capability)
            )