#!/usr/bin/env python3
"""
In-Memory Concurrent LFU Cache with Async TTL Eviction and Atomic Transactions.
Pure Python 3.11+ implementation. No external dependencies.
"""

import asyncio
import time
from dataclasses import dataclass, field
from typing import Any, Optional, Dict, Set, List


# ---------------------------------------------------------------------------
# Core Data Structures for O(1) Operations
# ---------------------------------------------------------------------------

@dataclass
class Node:
    key: Any
    value: Any
    freq: int
    expires_at: Optional[float]
    prev: Optional['Node'] = None
    next: Optional['Node'] = None


class DoublyLinkedList:
    """Sentinel-based doubly linked list supporting O(1) add/remove/pop."""
    def __init__(self):
        self.head = Node(None, None, 0, None)
        self.tail = Node(None, None, 0, None)
        self.head.next = self.tail
        self.tail.prev = self.head
        self.size = 0

    def add_to_tail(self, node: Node):
        node.prev = self.tail.prev
        node.next = self.tail
        self.tail.prev.next = node
        self.tail.prev = node
        self.size += 1

    def remove_node(self, node: Node):
        if node.prev is None or node.next is None:
            return
        node.prev.next = node.next
        node.next.prev = node.prev
        node.prev = node.next = None
        self.size -= 1

    def pop_head(self) -> Optional[Node]:
        if self.size == 0:
            return None
        node = self.head.next
        self.remove_node(node)
        return node

    def is_empty(self) -> bool:
        return self.size == 0


# ---------------------------------------------------------------------------
# Transaction Isolation Layer
# ---------------------------------------------------------------------------

class Transaction:
    """
    Provides ACID-like sub-sessions with Read-Your-Own-Writes isolation.
    Changes remain local until commit(), where they are applied atomically.
    """
    def __init__(self, cache: 'LFUCache'):
        self._cache = cache
        self._local_writes: Dict[Any, tuple] = {}      # key -> (value, ttl)
        self._local_deletes: Set[Any] = set()
        self._local_freq_bumps: Set[Any] = set()        # Track accesses for LFU bump
        self._committed = False
        self._rolled_back = False

    async def put(self, key: Any, value: Any, ttl_seconds: Optional[float] = None):
        if self._committed or self._rolled_back:
            raise RuntimeError("Transaction already finished")
        self._local_writes[key] = (value, ttl_seconds)
        self._local_deletes.discard(key)
        self._local_freq_bumps.add(key)

    async def get(self, key: Any) -> Optional[Any]:
        if self._committed or self._rolled_back:
            raise RuntimeError("Transaction already finished")
        # Read Your Own Writes
        if key in self._local_writes:
            val, _ = self._local_writes[key]
            self._local_freq_bumps.add(key)
            return val
        if key in self._local_deletes:
            return None
        # Read Global State
        val = await self._cache.get(key)
        if val is not None:
            self._local_freq_bumps.add(key)
        return val

    async def delete(self, key: Any) -> bool:
        if self._committed or self._rolled_back:
            raise RuntimeError("Transaction already finished")
        self._local_deletes.add(key)
        self._local_writes.pop(key, None)
        return True

    async def commit(self) -> bool:
        if self._committed or self._rolled_back:
            raise RuntimeError("Transaction already finished")
        self._committed = True

        async with self._cache._lock:
            # 1. Apply Writes
            for key, (value, ttl) in self._local_writes.items():
                if key in self._cache.store:
                    node = self._cache.store[key]
                    if self._cache._is_expired(node):
                        self._cache._remove_node(key)
                    else:
                        node.value = value
                        node.expires_at = time.time() + ttl if ttl else None
                        self._cache._update_freq(node)
                else:
                    self._cache._add_node(key, value, ttl)

            # 2. Apply Deletes
            for key in self._local_deletes:
                self._cache._remove_node(key)

            # 3. Apply Frequency Bumps (for reads during transaction)
            for key in self._local_freq_bumps:
                if key in self._cache.store and key not in self._local_deletes:
                    node = self._cache.store[key]
                    self._cache._update_freq(node)

            self._cache._cleanup_freq_lists()
        return True

    async def rollback(self):
        if self._committed or self._rolled_back:
            raise RuntimeError("Transaction already finished")
        self._rolled_back = True
        # Completely discard pending changes without mutating global state
        self._local_writes.clear()
        self._local_deletes.clear()
        self._local_freq_bumps.clear()


# ---------------------------------------------------------------------------
# Main Cache Implementation
# ---------------------------------------------------------------------------

class LFUCache:
    """
    Concurrent LFU Cache with O(1) get/put, dual-layer TTL eviction,
    and atomic transaction support.
    """
    def __init__(self, capacity: int, eviction_interval: float = 0.5):
        if capacity <= 0:
            raise ValueError("Cache capacity must be positive")
        self.capacity = capacity
        self.store: Dict[Any, Node] = {}
        self.freq_map: Dict[int, DoublyLinkedList] = {}
        self.min_freq: int = 0
        self.eviction_interval = eviction_interval
        self._lock = asyncio.Lock()
        self._evictor_task: Optional[asyncio.Task] = None
        self._evictor_running = False
        self._ttl_keys: Set[Any] = set()

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

    async def start_evictor(self):
        """Starts the background async TTL eviction loop."""
        if self._evictor_running:
            return
        self._evictor_running = True
        self._evictor_task = asyncio.create_task(self._eviction_loop())

    async def stop_evictor(self):
        """Stops the background async TTL eviction loop."""
        self._evictor_running = False
        if self._evictor_task:
            self._evictor_task.cancel()
            try:
                await self._evictor_task
            except asyncio.CancelledError:
                pass
            self._evictor_task = None

    async def _eviction_loop(self):
        while self._evictor_running:
            await asyncio.sleep(self.eviction_interval)
            if not self._evictor_running:
                break
            await self._evict_batch()

    async def _evict_batch(self):
        """Background sweep: acquires lock briefly, removes expired keys in batches."""
        async with self._lock:
            expired_keys = [
                k for k in list(self._ttl_keys)
                if k in self.store and self.store[k].expires_at and time.time() > self.store[k].expires_at
            ]
            batch = expired_keys[:100]  # Small batch to avoid lock contention
            for k in batch:
                self._remove_node(k)
                self._ttl_keys.discard(k)
            self._cleanup_freq_lists()

    def _cleanup_freq_lists(self):
        """Remove empty frequency buckets and update min_freq."""
        empty_freqs = [f for f, dll in self.freq_map.items() if dll.is_empty()]
        for f in empty_freqs:
            del self.freq_map[f]
        if self.freq_map:
            self.min_freq = min(self.freq_map.keys())
        else:
            self.min_freq = 0

    def _is_expired(self, node: Node) -> bool:
        return node.expires_at is not None and time.time() > node.expires_at

    def _remove_node(self, key: Any):
        """Removes a key from store and its frequency linked list."""
        if key not in self.store:
            return
        node = self.store.pop(key)
        self._ttl_keys.discard(key)
        freq = node.freq
        if freq in self.freq_map:
            self.freq_map[freq].remove_node(node)

    def _update_freq(self, node: Node):
        """Moves node to next frequency bucket in O(1)."""
        freq = node.freq
        self.freq_map[freq].remove_node(node)
        if self.freq_map[freq].is_empty():
            del self.freq_map[freq]
            if self.min_freq == freq:
                self.min_freq += 1
        node.freq += 1
        new_freq = node.freq
        if new_freq not in self.freq_map:
            self.freq_map[new_freq] = DoublyLinkedList()
        self.freq_map[new_freq].add_to_tail(node)

    async def get(self, key: Any) -> Optional[Any]:
        """O(1) get with lazy TTL eviction and frequency bump."""
        async with self._lock:
            if key not in self.store:
                return None
            node = self.store[key]
            if self._is_expired(node):
                self._remove_node(key)
                return None
            self._update_freq(node)
            return node.value

    async def put(self, key: Any, value: Any, ttl_seconds: Optional[float] = None):
        """O(1) put with capacity eviction and TTL handling."""
        async with self._lock:
            if key in self.store:
                node = self.store[key]
                if self._is_expired(node):
                    self._remove_node(key)
                else:
                    node.value = value
                    node.expires_at = time.time() + ttl_seconds if ttl_seconds else None
                    self._update_freq(node)
                    return
            if len(self.store) >= self.capacity:
                self._evict()
            self._add_node(key, value, ttl_seconds)

    def _evict(self):
        """Evicts the least frequently used item (O(1) via min_freq bucket)."""
        if self.min_freq in self.freq_map:
            dll = self.freq_map[self.min_freq]
            node = dll.pop_head()
            if node:
                self.store.pop(node.key, None)
                self._ttl_keys.discard(node.key)
                self._cleanup_freq_lists()

    def _add_node(self, key: Any, value: Any, ttl_seconds: Optional[float]):
        """Inserts new node into freq 1 bucket."""
        node = Node(key, value, 1, time.time() + ttl_seconds if ttl_seconds else None)
        self.store[key] = node
        if ttl_seconds:
            self._ttl_keys.add(key)
        if 1 not in self.freq_map:
            self.freq_map[1] = DoublyLinkedList()
            self.min_freq = 1
        self.freq_map[1].add_to_tail(node)

    async def delete(self, key: Any) -> bool:
        async with self._lock:
            if key in self.store:
                self._remove_node(key)
                return True
            return False

    async def size(self) -> int:
        async with self._lock:
            return len(self.store)


# ---------------------------------------------------------------------------
# Executable Unit Tests
# ---------------------------------------------------------------------------

async def test_lfu_eviction_order():
    """a) O(1) LFU eviction order when capacity is reached."""
    cache = LFUCache(3)
    await cache.put(1, 'a')
    await cache.put(2, 'b')
    await cache.put(3, 'c')
    
    # Access 1 twice to make it most frequent
    await cache.get(1)
    await cache.get(1)
    
    # Insert 4 -> should evict 2 (freq 1, least used)
    await cache.put(4, 'd')
    
    assert await cache.get(2) is None, "Failed: Least frequent key should be evicted"
    assert await cache.get(1) == 'a', "Failed: Key 1 should still exist"
    assert await cache.get(3) == 'c', "Failed: Key 3 should still exist"
    assert await cache.get(4) == 'd', "Failed: Key 4 should exist"
    print("✅ Test a) LFU eviction order: PASSED")


async def test_ttl_eviction():
    """b) Lazy TTL vs. Background Async Sweep eviction."""
    cache = LFUCache(10)
    await cache.start_evictor()
    
    # Lazy eviction
    await cache.put('lazy', 'val', ttl_seconds=0.1)
    assert await cache.get('lazy') == 'val', "Failed: Should return value before TTL"
    await asyncio.sleep(0.15)
    assert await cache.get('lazy') is None, "Failed: Lazy eviction should trigger on next get"
    
    # Background sweep eviction
    await cache.put('bg', 'val2', ttl_seconds=0.1)
    await asyncio.sleep(0.25)  # Wait for background task cycle
    assert await cache.get('bg') is None, "Failed: Background sweep should evict expired key"
    
    await cache.stop_evictor()
    print("✅ Test b) TTL eviction (Lazy & Background): PASSED")


async def test_transaction_isolation():
    """c) Transaction commit visibility vs. rollback state restoration."""
    cache = LFUCache(10)
    await cache.put('a', 1)
    
    # Commit test
    tx = cache.begin_transaction()
    await tx.put('a', 2)
    await tx.put('b', 3)
    
    # Read your own writes
    assert await tx.get('a') == 2, "Failed: Should see local write"
    assert await tx.get('b') == 3, "Failed: Should see local write"
    
    # Global shouldn't see uncommitted changes
    assert await cache.get('a') == 1, "Failed: Global should not see uncommitted write"
    
    await tx.commit()
    assert await cache.get('a') == 2, "Failed: Global should see committed write"
    assert await cache.get('b') == 3, "Failed: Global should see committed write"
    
    # Rollback test
    tx2 = cache.begin_transaction()
    await tx2.put('c', 4)
    assert await tx2.get('c') == 4, "Failed: Should see local write"
    await tx2.rollback()
    assert await cache.get('c') is None, "Failed: Rollback should discard changes"
    print("✅ Test c) Transaction isolation (Commit & Rollback): PASSED")


async def test_stress_concurrency():
    """d) Stress test with 50 concurrent async tasks reading/writing simultaneously."""
    cache = LFUCache(100)
    await cache.start_evictor()
    
    tasks = []
    for i in range(50):
        async def worker(idx: int):
            for j in range(10):
                key = f"k{idx}"
                await cache.put(key, f"v{idx}_{j}")
                val = await cache.get(key)
                assert val == f"v{idx}_{j}", f"Failed: Read mismatch for {key}"
        tasks.append(asyncio.create_task(worker(i)))
        
    await asyncio.gather(*tasks)
    
    # All 50 keys should exist
    assert await cache.size() == 50, f"Failed: Expected 50 keys, got {await cache.size()}"
    
    # Verify concurrent reads/writes didn't corrupt data
    for i in range(50):
        assert await cache.get(f"k{i}") == f"v{i}_9", f"Failed: Data corruption for k{i}"
        
    await cache.stop_evictor()
    print("✅ Test d) Stress concurrency (50 tasks): PASSED")


async def main():
    print("🚀 Running LFU Cache Unit Tests...\n")
    await test_lfu_eviction_order()
    await test_ttl_eviction()
    await test_transaction_isolation()
    await test_stress_concurrency()
    print("\n🎉 All tests passed successfully.")


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