#!/usr/bin/env python3
"""
In-Memory Concurrent LFU Cache with Async TTL Eviction & Atomic Transactions.

Pure Python 3.11+, no external dependencies.

O(1) average-time get/put via frequency buckets with doubly-linked lists.
Dual-layer TTL: lazy on access + background async sweep.
Atomic transactions with read-your-own-writes, commit, and rollback.
"""

import asyncio
import time
from typing import (Dict, Generic, Hashable, Optional, Set, TypeVar)


KT = TypeVar('KT', bound=Hashable)
VT = TypeVar('VT')

_SENTINEL = object()


class LFUCacheError(Exception):
    """Base exception for LFU cache errors."""


class _Node(Generic[KT, VT]):
    """Doubly-linked list node holding key, value, frequency, and expiry."""

    __slots__ = ('key', 'value', 'freq', 'expires_at', 'prev', 'next')

    def __init__(self, key: KT, value: VT) -> None:
        self.key = key
        self.value = value
        self.freq: int = 0
        self.expires_at: Optional[float] = None
        self.prev: Optional['_Node[KT, VT]'] = None
        self.next: Optional['_Node[KT, VT]'] = None


class _DLL(Generic[KT, VT]):
    """
    Doubly-linked list with sentinel head/tail.
    O(1) append (tail/MRU), remove (by reference), and pop_left (head/LRU).
    """

    __slots__ = ('_head', '_tail', 'size')

    def __init__(self) -> None:
        self._head = _Node[KT, VT](_SENTINEL, None)  # type: ignore[arg-type]
        self._tail = _Node[KT, VT](_SENTINEL, None)  # type: ignore[arg-type]
        self._head.next = self._tail
        self._tail.prev = self._head
        self.size: int = 0

    def append(self, node: _Node[KT, VT]) -> None:
        """Add node to tail (most-recently-used position). O(1)."""
        last = self._tail.prev
        last.next = node
        node.prev = last
        node.next = self._tail
        self._tail.prev = node
        self.size += 1

    def remove(self, node: _Node[KT, VT]) -> None:
        """Remove node that must be in this list. O(1)."""
        node.prev.next = node.next
        node.next.prev = node.prev
        node.prev = None
        node.next = None
        self.size -= 1

    def pop_left(self) -> Optional[_Node[KT, VT]]:
        """Remove and return the head node (LRU position). O(1)."""
        if self.size == 0:
            return None
        node = self._head.next
        self.remove(node)
        return node


class Transaction(Generic[KT, VT]):
    """
    Atomic transaction with read-your-own-writes isolation.

    Uncommitted writes are invisible to global cache readers.
    All local state is discarded on rollback without touching global structures.
    """

    def __init__(self, cache: 'LFUCache[KT, VT]') -> None:
        self._cache = cache
        self._writes: Dict[KT, VT] = {}
        self._deletes: Set[KT] = set()
        self._snapshot: Dict[KT, VT] = {}
        self._freq_deltas: Dict[KT, int] = {}
        self._ttls: Dict[KT, Optional[float]] = {}
        self._committed: bool = False
        self._rolled_back: bool = False

    def _check_active(self) -> None:
        if self._committed:
            raise LFUCacheError("Transaction already committed")
        if self._rolled_back:
            raise LFUCacheError("Transaction already rolled back")

    async def get(self, key: KT) -> Optional[VT]:
        """Read key with read-your-own-writes and snapshot isolation."""
        self._check_active()

        if key in self._deletes:
            return None
        if key in self._writes:
            self._freq_deltas[key] = self._freq_deltas.get(key, 0) + 1
            return self._writes[key]
        if key in self._snapshot:
            self._freq_deltas[key] = self._freq_deltas.get(key, 0) + 1
            return self._snapshot[key]

        async with self._cache._lock:
            node = self._cache._key_to_node.get(key)
            if node is None:
                return None
            if node.expires_at is not None and node.expires_at <= time.monotonic():
                self._cache._remove_key(key)
                return None
            self._snapshot[key] = node.value

        self._freq_deltas[key] = self._freq_deltas.get(key, 0) + 1
        return self._snapshot[key]

    async def put(self, key: KT, value: VT,
                  ttl_seconds: Optional[float] = None) -> None:
        """Write key within transaction."""
        self._check_active()
        self._writes[key] = value
        self._deletes.discard(key)
        if ttl_seconds is not None:
            self._ttls[key] = ttl_seconds if ttl_seconds > 0 else None
        self._freq_deltas[key] = self._freq_deltas.get(key, 0) + 1

    async def delete(self, key: KT) -> None:
        """Mark key for deletion on commit."""
        self._check_active()
        self._deletes.add(key)
        self._writes.pop(key, None)
        self._snapshot.pop(key, None)
        self._freq_deltas.pop(key, None)
        self._ttls.pop(key, None)

    async def commit(self) -> None:
        """Atomically apply all pending changes to the global cache."""
        if self._rolled_back:
            raise LFUCacheError("Transaction already rolled back")
        if self._committed:
            return

        async with self._cache._lock:
            try:
                for key in self._deletes:
                    self._cache._remove_key(key)

                for key, value in self._writes.items():
                    node = self._cache._key_to_node.get(key)
                    ttl = self._ttls.get(key)
                    delta = self._freq_deltas.get(key, 1)

                    if node is not None:
                        node.value = value
                        if ttl is not None:
                            node.expires_at = (
                                time.monotonic() + ttl if ttl is not None else None
                            )
                        self._cache._change_freq(node, delta)
                    else:
                        while len(self._cache._key_to_node) >= self._cache._capacity:
                            self._cache._evict_one()
                        node = _Node(key, value)
                        node.freq = delta
                        if ttl is not None:
                            node.expires_at = time.monotonic() + ttl
                        self._cache._key_to_node[key] = node
                        self._cache._add_to_freq_list(node)

                for key, delta in self._freq_deltas.items():
                    if key not in self._writes and key not in self._deletes:
                        node = self._cache._key_to_node.get(key)
                        if node is not None:
                            self._cache._change_freq(node, delta)

                while len(self._cache._key_to_node) > self._cache._capacity:
                    self._cache._evict_one()

                self._committed = True
            except BaseException:
                self._rolled_back = True
                raise

    async def rollback(self) -> None:
        """Discard all pending changes. Global cache is untouched."""
        if self._committed:
            raise LFUCacheError("Transaction already committed")
        self._rolled_back = True


class LFUCache(Generic[KT, VT]):
    """
    In-Memory Concurrent LFU Cache with Async TTL Eviction.

    O(1) average-time get/put using frequency buckets with doubly-linked lists.
    Dual-layer TTL eviction: lazy on access + background async sweep.
    """

    def __init__(self, capacity: int = 1000,
                 evictor_interval: float = 1.0,
                 evictor_batch_size: int = 10,
                 evictor_scan_budget: int = 100) -> None:
        if capacity < 1:
            raise ValueError("Capacity must be >= 1")

        self._capacity = capacity
        self._evictor_interval = evictor_interval
        self._evictor_batch_size = evictor_batch_size
        self._evictor_scan_budget = evictor_scan_budget

        self._key_to_node: Dict[KT, _Node[KT, VT]] = {}
        self._freq_to_list: Dict[int, _DLL[KT, VT]] = {}
        self._min_freq: int = 0
        self._lock = asyncio.Lock()
        self._evictor_running: bool = False
        self._evictor_task: Optional[asyncio.Task] = None

    # ---- Public API ---------------------------------------------------------

    async def get(self, key: KT) -> Optional[VT]:
        """Retrieve value by key. Returns None if missing or expired.

        Accesses increment the key's frequency (LFU tracking).
        Expired keys are lazily evicted on access.
        """
        async with self._lock:
            node = self._key_to_node.get(key)
            if node is None:
                return None
            if node.expires_at is not None and node.expires_at <= time.monotonic():
                self._remove_key(key)
                return None
            self._change_freq(node, 1)
            return node.value

    async def put(self, key: KT, value: VT,
                  ttl_seconds: Optional[float] = None) -> None:
        """Insert or update a key-value pair.

        If *ttl_seconds* is None (default) the entry never expires.
        If *ttl_seconds* is <= 0 it is treated as no expiration.
        If the cache is at capacity the least-frequently-used item is evicted
        (tie-broken by least-recently-used within the minimum frequency tier).
        """
        async with self._lock:
            node = self._key_to_node.get(key)

            if node is not None:
                if node.expires_at is not None and node.expires_at <= time.monotonic():
                    self._remove_key(key)
                    node = None
                else:
                    node.value = value
                    if ttl_seconds is not None:
                        node.expires_at = (
                            time.monotonic() + ttl_seconds if ttl_seconds > 0 else None
                        )
                    self._change_freq(node, 1)
                    return

            while len(self._key_to_node) >= self._capacity:
                self._evict_one()

            node = _Node(key, value)
            node.freq = 1
            if ttl_seconds is not None and ttl_seconds > 0:
                node.expires_at = time.monotonic() + ttl_seconds
            self._key_to_node[key] = node
            self._add_to_freq_list(node)

    async def delete(self, key: KT) -> bool:
        """Remove *key* from the cache. Returns True if the key existed."""
        async with self._lock:
            return self._remove_key(key)

    def begin_transaction(self) -> Transaction[KT, VT]:
        """Open an atomic transaction for batched reads/writes."""
        return Transaction(self)

    @property
    def capacity(self) -> int:
        return self._capacity

    @property
    def size(self) -> int:
        return len(self._key_to_node)

    # ---- Background TTL Evictor --------------------------------------------

    def start_evictor(self) -> None:
        """Launch the background TTL eviction loop as an asyncio task."""
        if self._evictor_running:
            return
        self._evictor_running = True
        self._evictor_task = asyncio.create_task(self._evictor_loop())

    async def stop_evictor(self) -> None:
        """Cancel and wait for the background eviction task to finish."""
        if not self._evictor_running:
            return
        self._evictor_running = False
        if self._evictor_task is not None:
            self._evictor_task.cancel()
            try:
                await self._evictor_task
            except asyncio.CancelledError:
                pass
            self._evictor_task = None

    async def _evictor_loop(self) -> None:
        """
        Background loop: periodically scan a limited batch of keys and remove
        expired entries.  Releases the lock between batches so concurrent
        reads/writes are not blocked for extended periods.
        """
        cursor = 0
        keys: list = []

        try:
            while self._evictor_running:
                await asyncio.sleep(self._evictor_interval)

                async with self._lock:
                    if cursor >= len(keys):
                        keys = list(self._key_to_node.keys())
                        cursor = 0
                    if not keys:
                        continue

                    now = time.monotonic()
                    removed = 0
                    scan = self._evictor_scan_budget

                    while scan > 0 and cursor < len(keys):
                        key = keys[cursor]
                        cursor += 1
                        scan -= 1

                        node = self._key_to_node.get(key)
                        if (node is not None and
                                node.expires_at is not None and
                                node.expires_at <= now):
                            self._remove_key(key)
                            removed += 1
                            if removed >= self._evictor_batch_size:
                                break
        except asyncio.CancelledError:
            pass

    # ---- O(1) Internal Helpers ---------------------------------------------

    def _add_to_freq_list(self, node: _Node[KT, VT]) -> None:
        """Insert *node* into its frequency bucket and update *min_freq*."""
        freq = node.freq
        if freq not in self._freq_to_list:
            self._freq_to_list[freq] = _DLL()
        self._freq_to_list[freq].append(node)
        if self._min_freq not in self._freq_to_list or freq < self._min_freq:
            self._min_freq = freq

    def _remove_from_freq_list(self, node: _Node[KT, VT]) -> None:
        """Remove *node* from its frequency bucket. No-op if not linked."""
        if node.prev is None or node.next is None:
            return
        freq = node.freq
        lst = self._freq_to_list.get(freq)
        if lst is None:
            return
        lst.remove(node)
        if lst.size == 0:
            del self._freq_to_list[freq]

    def _change_freq(self, node: _Node[KT, VT], delta: int) -> None:
        """Atomically increase *node*'s frequency by *delta* and
        move it to the corresponding bucket.  O(1)."""
        if delta <= 0:
            return
        self._remove_from_freq_list(node)
        node.freq += delta
        self._add_to_freq_list(node)

    def _remove_key(self, key: KT) -> bool:
        """Remove *key* from all internal structures. Returns True if existed."""
        node = self._key_to_node.pop(key, None)
        if node is None:
            return False
        self._remove_from_freq_list(node)
        return True

    def _evict_one(self) -> Optional[KT]:
        """
        Evict one item: LRU among the minimum-frequency bucket.
        Returns the evicted key, or None if the cache is empty.
        """
        if not self._key_to_node:
            return None

        if self._min_freq not in self._freq_to_list:
            if not self._freq_to_list:
                return None
            self._min_freq = min(self._freq_to_list)

        lst = self._freq_to_list.get(self._min_freq)
        if lst is None or lst.size == 0:
            return None

        node = lst.pop_left()
        if node is None:
            return None

        del self._key_to_node[node.key]
        if lst.size == 0:
            del self._freq_to_list[self._min_freq]
            if self._freq_to_list:
                self._min_freq = min(self._freq_to_list)

        return node.key

    # ---- Async Context Manager ---------------------------------------------

    async def __aenter__(self) -> 'LFUCache[KT, VT]':
        return self

    async def __aexit__(self, *exc_info) -> None:
        await self.stop_evictor()


# ============================================================================
#  Unit Tests
# ============================================================================

async def _run_tests() -> None:
    passed = 0
    total = 0

    def check(cond: bool, msg: str):
        nonlocal passed, total
        total += 1
        if cond:
            passed += 1
        else:
            print(f"  FAIL: {msg}")

    # ------------------------------------------------------------------
    # 1. O(1) LFU eviction order
    # ------------------------------------------------------------------
    print("1. LFU eviction order ... ", end="", flush=True)

    cache = LFUCache(capacity=3)

    await cache.put("a", 1)
    await cache.put("b", 2)
    await cache.put("c", 3)

    await cache.get("a")
    await cache.get("a")
    await cache.get("b")
    # freqs: a=3, b=2, c=1

    await cache.put("d", 4)        # evicts c (freq=1)
    check(await cache.get("c") is None, "c was evicted (freq=1)")
    check(await cache.get("d") == 4, "d is present")
    check(await cache.get("a") == 1, "a is present")
    check(await cache.get("b") == 2, "b is present")

    # a=4, b=3, d=1 after the gets above
    await cache.put("e", 5)        # evicts d (freq=1)
    check(await cache.get("d") is None, "d was evicted")
    check(await cache.get("e") == 5, "e is present")
    check(cache.size == 3, "cache size == 3")

    print("OK")

    # ------------------------------------------------------------------
    # 2. TTL eviction
    # ------------------------------------------------------------------
    print("2. TTL eviction ... ", end="", flush=True)

    # 2a. Lazy eviction
    cache2 = LFUCache(capacity=10)
    await cache2.put("lazy", "alive", ttl_seconds=0.02)
    await asyncio.sleep(0.03)
    check(await cache2.get("lazy") is None, "lazy TTL eviction on get")

    await cache2.put("lazy2", "alive", ttl_seconds=0.02)
    await asyncio.sleep(0.03)
    check(await cache2.get("lazy2") is None, "lazy TTL eviction on second get")

    # 2b. Background eviction
    cache3 = LFUCache(capacity=10, evictor_interval=0.02, evictor_batch_size=5)
    cache3.start_evictor()
    await cache3.put("bg", "data", ttl_seconds=0.01)
    await asyncio.sleep(0.10)
    check(await cache3.get("bg") is None, "background evictor removes expired key")
    await cache3.stop_evictor()

    print("OK")

    # ------------------------------------------------------------------
    # 3. Transactions
    # ------------------------------------------------------------------
    print("3. Transactions ... ", end="", flush=True)

    # 3a. Commit
    cache4 = LFUCache(capacity=10)
    await cache4.put("x", 10)

    tx = cache4.begin_transaction()
    await tx.put("x", 20)
    check(await tx.get("x") == 20, "tx reads its own write")
    check(await cache4.get("x") == 10, "global does NOT see uncommitted write")
    await tx.commit()
    check(await cache4.get("x") == 20, "global sees committed value")

    # 3b. Rollback
    await cache4.put("y", 100)
    tx2 = cache4.begin_transaction()
    await tx2.put("y", 200)
    await tx2.rollback()
    check(await cache4.get("y") == 100, "rollback discards writes")

    # 3c. Delete in transaction
    await cache4.put("z", 300)
    tx3 = cache4.begin_transaction()
    await tx3.delete("z")
    check(await tx3.get("z") is None, "tx sees its own delete")
    check(await cache4.get("z") == 300, "global not affected before commit")
    await tx3.commit()
    check(await cache4.get("z") is None, "global sees committed delete")

    # 3d. Rollback after delete
    await cache4.put("w", 400)
    tx4 = cache4.begin_transaction()
    await tx4.delete("w")
    await tx4.rollback()
    check(await cache4.get("w") == 400, "rollback restores deleted key")

    # 3e. Read own writes after rollback raises
    tx5 = cache4.begin_transaction()
    await tx5.put("n", 1)
    await tx5.rollback()
    try:
        await tx5.get("n")
        check(False, "get on rolled-back tx should raise")
    except LFUCacheError:
        check(True, "rolled-back tx raises on get")

    print("OK")

    # ------------------------------------------------------------------
    # 4. Concurrent stress
    # ------------------------------------------------------------------
    print("4. Concurrent stress (50 workers) ... ", end="", flush=True)

    cache5 = LFUCache(capacity=50, evictor_interval=0.05, evictor_batch_size=5)
    cache5.start_evictor()

    async def worker(uid: int):
        for i in range(30):
            key = f"k_{(uid + i) % 40}"
            await cache5.put(key, uid * 1000 + i, ttl_seconds=0.5)
            _ = await cache5.get(key)
            await cache5.get(f"k_{(uid + i + 1) % 40}")
            if i % 5 == 0:
                await cache5.delete(f"k_{(uid + i + 2) % 40}")
            await asyncio.sleep(0)

    workers = [worker(i) for i in range(50)]
    await asyncio.gather(*workers)
    await cache5.stop_evictor()

    check(cache5.size <= 50, f"cache size {cache5.size} <= capacity 50")
    print("OK")

    # ------------------------------------------------------------------
    print(f"\nResults: {passed}/{total} checks passed")
    if passed < total:
        raise SystemExit(1)


async def main() -> None:
    await _run_tests()


if __name__ == "__main__":
    asyncio.run(main())