# lfu_cache.py
import asyncio
import time
from typing import Any, Optional, Dict, Set

class _Node:
    __slots__ = ("key", "value", "expire_at", "freq", "prev", "next")
    def __init__(self, key=None, value=None, expire_at=None, freq=0):
        self.key = key
        self.value = value
        self.expire_at = expire_at
        self.freq = freq
        self.prev = None
        self.next = None


class _FreqList:
    __slots__ = ("head", "tail")
    def __init__(self):
        self.head = _Node()
        self.tail = _Node()
        self.head.next = self.tail
        self.tail.prev = self.head


class LFUCache:
    def __init__(self, capacity: int, evict_interval: float = 0.5):
        self.capacity = capacity
        self.key_map: Dict[Any, _Node] = {}
        self.freq_map: Dict[int, _FreqList] = {}
        self.min_freq = 0
        self.lock = asyncio.Lock()
        self._evictor_task: Optional[asyncio.Task] = None
        self._evict_interval = evict_interval

    # ---------- internal helpers ----------
    def _update_min_freq(self):
        while self.min_freq not in self.freq_map:
            self.min_freq += 1

    def _add_node_to_head(self, freq: int, node: _Node):
        lst = self.freq_map.get(freq)
        if lst is None:
            lst = _FreqList()
            self.freq_map[freq] = lst
        node.prev = lst.head
        node.next = lst.head.next
        lst.head.next.prev = node
        lst.head.next = node

    def _remove_node(self, node: _Node):
        freq = node.freq
        lst = self.freq_map.get(freq)
        if lst:
            node.prev.next = node.next
            node.next.prev = node.prev
            if lst.head.next is lst.tail:
                del self.freq_map[freq]
                if self.min_freq == freq:
                    self._update_min_freq()
        if self.key_map.get(node.key) is node:
            del self.key_map[node.key]

    def _increment_freq(self, node: _Node):
        old_freq = node.freq
        lst = self.freq_map[old_freq]
        node.prev.next = node.next
        node.next.prev = node.prev
        if lst.head.next is lst.tail:
            del self.freq_map[old_freq]
            if self.min_freq == old_freq:
                self._update_min_freq()
        new_freq = old_freq + 1
        node.freq = new_freq
        self._add_node_to_head(new_freq, node)

    def _evict_one(self):
        if not self.freq_map:
            return
        if self.min_freq not in self.freq_map:
            self._update_min_freq()
        lst = self.freq_map[self.min_freq]
        node = lst.tail.prev
        if node is lst.head:
            return
        self._remove_node(node)

    def _is_expired(self, node: _Node) -> bool:
        return node.expire_at is not None and time.monotonic() > node.expire_at

    # ---------- public API ----------
    async def get(self, key: Any) -> Optional[Any]:
        async with self.lock:
            node = self.key_map.get(key)
            if not node:
                return None
            if self._is_expired(node):
                self._remove_node(node)
                return None
            self._increment_freq(node)
            return node.value

    async def put(self, key: Any, value: Any, ttl_seconds: Optional[float] = None):
        async with self.lock:
            now = time.monotonic()
            expire_at = now + ttl_seconds if ttl_seconds is not None else None
            node = self.key_map.get(key)
            if node:
                if self._is_expired(node):
                    self._remove_node(node)
                    node = None
            if node:
                node.value = value
                node.expire_at = expire_at
                self._increment_freq(node)
                return
            if len(self.key_map) >= self.capacity:
                self._evict_one()
            node = _Node(key=key, value=value, expire_at=expire_at, freq=1)
            self.key_map[key] = node
            self._add_node_to_head(1, node)
            if self.min_freq == 0 or 1 < self.min_freq:
                self.min_freq = 1

    def begin_transaction(self) -> "Transaction":
        return Transaction(self)

    async def start_evictor(self):
        if self._evictor_task and not self._evictor_task.done():
            return
        self._evictor_task = asyncio.create_task(self._evictor_loop())

    async def stop_evictor(self):
        if self._evictor_task:
            self._evictor_task.cancel()
            try:
                await self._evictor_task
            except asyncio.CancelledError:
                pass
            self._evictor_task = None

    async def _evictor_loop(self):
        while True:
            await asyncio.sleep(self._evict_interval)
            async with self.lock:
                now = time.monotonic()
                batch = 100
                count = 0
                for k in list(self.key_map.keys()):
                    if count >= batch:
                        break
                    node = self.key_map.get(k)
                    if node and node.expire_at and now > node.expire_at:
                        self._remove_node(node)
                        count += 1


class Transaction:
    def __init__(self, cache: LFUCache):
        self.cache = cache
        self.pending_puts: Dict[Any, tuple[Any, Optional[float]]] = {}
        self.pending_deletes: Set[Any] = set()

    async def get(self, key: Any) -> Optional[Any]:
        if key in self.pending_deletes:
            return None
        if key in self.pending_puts:
            val, _ = self.pending_puts[key]
            return val
        async with self.cache.lock:
            node = self.cache.key_map.get(key)
            if not node:
                return None
            if node.expire_at and time.monotonic() > node.expire_at:
                self.cache._remove_node(node)
                return None
            return node.value

    async def put(self, key: Any, value: Any, ttl_seconds: Optional[float] = None):
        self.pending_puts[key] = (value, ttl_seconds)
        self.pending_deletes.discard(key)

    async def delete(self, key: Any):
        self.pending_deletes.add(key)
        self.pending_puts.pop(key, None)

    async def commit(self):
        async with self.cache.lock:
            for key in self.pending_deletes:
                node = self.cache.key_map.get(key)
                if node:
                    self.cache._remove_node(node)
            for key, (value, ttl_seconds) in self.pending_puts.items():
                expire_at = None
                if ttl_seconds is not None:
                    expire_at = time.monotonic() + ttl_seconds
                node = self.cache.key_map.get(key)
                if node:
                    node.value = value
                    node.expire_at = expire_at
                    self.cache._increment_freq(node)
                else:
                    if len(self.cache.key_map) >= self.cache.capacity:
                        self.cache._evict_one()
                    node = _Node(key=key, value=value, expire_at=expire_at, freq=1)
                    self.cache.key_map[key] = node
                    self.cache._add_node_to_head(1, node)
                    if self.cache.min_freq == 0 or 1 < self.cache.min_freq:
                        self.cache.min_freq = 1
            self.pending_puts.clear()
            self.pending_deletes.clear()

    def rollback(self):
        self.pending_puts.clear()
        self.pending_deletes.clear()


# ------------------- tests -------------------
async def main():
    cache = LFUCache(capacity=3, evict_interval=0.2)
    await cache.start_evictor()
    print("=== LFU eviction order ===")
    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")
    await cache.put("d", 4)  # should evict c
    assert await cache.get("c") is None
    assert await cache.get("a") == 1
    print("LFU eviction OK")

    print("=== TTL lazy vs background ===")
    await cache.put("x", "temp", ttl_seconds=0.1)
    await asyncio.sleep(0.15)
    assert await cache.get("x") is None
    await cache.put("y", "temp2", ttl_seconds=0.1)
    await asyncio.sleep(0.3)
    assert await cache.get("y") is None
    print("TTL OK")

    print("=== Transaction commit / rollback ===")
    tx = cache.begin_transaction()
    await tx.put("t1", 100)
    assert await cache.get("t1") is None
    await tx.commit()
    assert await cache.get("t1") == 100

    tx2 = cache.begin_transaction()
    await tx2.put("t2", 200)
    await tx2.rollback()
    assert await cache.get("t2") is None
    print("Transaction OK")

    print("=== Stress test 50 concurrent tasks ===")
    async def worker(i):
        for j in range(20):
            k = f"k{i%10}"
            await cache.put(k, i*j, ttl_seconds=5)
            await cache.get(k)
    await asyncio.gather(*(worker(i) for i in range(50)))
    print("Stress OK")

    await cache.stop_evictor()
    print("All tests passed")

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