"""
In-Memory Concurrent LFU Cache with Async TTL Eviction & Atomic Transactions.
Pure Python 3.11+ implementation using asyncio, dataclasses, and collections.OrderedDict.
"""

import asyncio
import time
from dataclasses import dataclass
from typing import Any, Dict, Optional, Set, Tuple
from collections import OrderedDict


@dataclass
class _Node:
    """Internal node storing cache value and TTL expiry timestamp."""
    key: str
    value: Any
    ttl_expiry: float


class LFUCache:
    """
    O(1) Least Frequently Used Cache with dual-layer TTL eviction and async concurrency.
    
    Data Structures:
      - self.cache: Dict[str, _Node] -> Direct O(1) key lookup
      - self.freq_map: Dict[int, OrderedDict[str, None]] -> Frequency buckets maintaining insertion order (LRU within same freq)
      - self.key_freq: Dict[str, int] -> Tracks current frequency of each key for O(1) updates
      - self.ttl_map: Dict[str, float] -> Stores absolute TTL expiry timestamps
    """

    def __init__(self, capacity: int):
        self.capacity = max(0, capacity)
        self.cache: Dict[str, _Node] = {}
        self.freq_map: Dict[int, OrderedDict[str, None]] = {}
        self.key_freq: Dict[str, int] = {}
        self.ttl_map: Dict[str, float] = {}
        self.min_freq: int = 0
        self._lock = asyncio.Lock()
        self._evictor_task: Optional[asyncio.Task] = None

    # ------------------------------------------------------------------ #
    #                          PUBLIC API                                  #
    # ------------------------------------------------------------------ #

    async def start_evictor(self) -> None:
        """Start the non-blocking background TTL eviction loop."""
        if self._evictor_task is not None:
            return
        self._evictor_task = asyncio.create_task(self._background_loop())

    async def stop_evictor(self) -> None:
        """Gracefully stop the background eviction task."""
        if self._evictor_task:
            self._evictor_task.cancel()
            try:
                await self._evictor_task
            except asyncio.CancelledError:
                pass
            self._evictor_task = None

    async def get(self, key: str) -> Optional[Any]:
        """
        Retrieve value by key. O(1) average time complexity.
        Performs lazy TTL eviction upon access.
        """
        current_time = time.monotonic()
        expiry = self.ttl_map.get(key)
        
        # Lazy TTL Eviction
        if expiry is not None and current_time >= expiry:
            await self._remove_key(key)
            return None

        async with self._lock:
            if key not in self.cache:
                return None
            
            node = self.cache[key]
            
            # Frequency Increment & Bucket Migration (O(1))
            old_freq = self.key_freq[key]
            new_freq = old_freq + 1
            self.key_freq[key] = new_freq

            old_bucket = self.freq_map[old_freq]
            del old_bucket[key]
            if not old_bucket:
                del self.freq_map[old_freq]
                if self.min_freq == old_freq:
                    self.min_freq = min(self.freq_map.keys()) if self.freq_map else 0

            new_bucket = self.freq_map.setdefault(new_freq, OrderedDict())
            new_bucket[key] = None  # Store key reference in bucket
            
            return node.value

    async def put(self, key: str, value: Any, ttl_seconds: float) -> None:
        """
        Insert or update key-value pair with TTL. O(1) average time complexity.
        Performs lazy TTL eviction before insertion if needed.
        """
        current_time = time.monotonic()
        expiry = self.ttl_map.get(key)
        
        # Lazy TTL Eviction for stale keys
        if expiry is not None and current_time >= expiry:
            await self._remove_key(key)

        async with self._lock:
            if key in self.cache:
                # Update existing: increment frequency & migrate bucket
                old_freq = self.key_freq[key]
                new_freq = old_freq + 1
                self.key_freq[key] = new_freq

                old_bucket = self.freq_map[old_freq]
                del old_bucket[key]
                if not old_bucket:
                    del self.freq_map[old_freq]
                    if self.min_freq == old_freq:
                        self.min_freq = min(self.freq_map.keys()) if self.freq_map else 0

                new_bucket = self.freq_map.setdefault(new_freq, OrderedDict())
                new_bucket[key] = None
            else:
                # Insert new: evict LFU if at capacity
                if len(self.cache) >= self.capacity and self.capacity > 0:
                    await self._evict_lfu()

                freq = 1
                self.key_freq[key] = freq
                bucket = self.freq_map.setdefault(freq, OrderedDict())
                bucket[key] = None
                self.min_freq = 1

            # Update node & TTL map
            self.cache[key] = _Node(key=key, value=value, ttl_expiry=current_time + ttl_seconds)
            self.ttl_map[key] = current_time + ttl_seconds

    async def delete(self, key: str) -> bool:
        """Delete a key from the cache. O(1)."""
        async with self._lock:
            if key not in self.cache:
                return False
            await self._remove_key(key)
            return True

    # ------------------------------------------------------------------ #
    #                          INTERNAL HELPERS                            #
    # ------------------------------------------------------------------ #

    async def _evict_lfu(self) -> None:
        """Evict the least frequently used key (oldest among ties). O(1)."""
        if not self.freq_map or self.min_freq == 0:
            return
        
        # Ensure min_freq points to a valid tier
        while self.min_freq in self.freq_map and self.min_freq < max(self.freq_map.keys()):
            pass
        if self.min_freq not in self.freq_map:
            self.min_freq = min(self.freq_map.keys()) if self.freq_map else 0

        bucket = self.freq_map[self.min_freq]
        evict_key, _ = bucket.popitem(last=False)
        await self._remove_key(evict_key)

    async def _remove_key(self, key: str) -> None:
        """Remove key from all internal structures. O(1)."""
        if key not in self.cache:
            return
            
        node = self.cache.pop(key)
        self.ttl_map.pop(key, None)
        
        freq = self.key_freq.pop(key)
        bucket = self.freq_map[freq]
        del bucket[key]
        
        if not bucket:
            del self.freq_map[freq]
            if self.min_freq == freq:
                # Find next valid minimum frequency
                self.min_freq = min(self.freq_map.keys()) if self.freq_map else 0

    async def _background_loop(self) -> None:
        """Non-blocking background task that purges expired keys in batches."""
        while True:
            await asyncio.sleep(0.1)  # Check interval
            current_time = time.monotonic()
            
            async with self._lock:
                # Collect expired keys safely
                expired_keys = [k for k, exp in list(self.ttl_map.items()) if current_time >= exp]
                
                # Purge in batch (yields control between removals implicitly via await)
                for key in expired_keys:
                    await self._remove_key(key)

    async def apply_transaction(self, tx: "Transaction") -> None:
        """Atomically apply transaction buffers to global state."""
        async with self._lock:
            # 1. Apply local deletes first (frees capacity for puts)
            for key in list(tx._local_deletes):
                if key in self.cache:
                    await self._remove_key(key)
                tx._local_puts.pop(key, None)

            # 2. Apply local puts
            for key, (value, expiry) in tx._local_puts.items():
                if key in self.cache:
                    old_freq = self.key_freq[key]
                    new_freq = old_freq + 1
                    self.key_freq[key] = new_freq

                    old_bucket = self.freq_map[old_freq]
                    del old_bucket[key]
                    if not old_bucket:
                        del self.freq_map[old_freq]
                        if self.min_freq == old_freq:
                            self.min_freq = min(self.freq_map.keys()) if self.freq_map else 0

                    new_bucket = self.freq_map.setdefault(new_freq, OrderedDict())
                    new_bucket[key] = None
                else:
                    if len(self.cache) >= self.capacity and self.capacity > 0:
                        await self._evict_lfu()
                    
                    freq = 1
                    self.key_freq[key] = freq
                    bucket = self.freq_map.setdefault(freq, OrderedDict())
                    bucket[key] = None
                    self.min_freq = 1

                self.cache[key] = _Node(key=key, value=value, ttl_expiry=expiry)
                self.ttl_map[key] = expiry


class Transaction:
    """
    ACID-like sub-session handle supporting Read-Your-Own-Writes and isolation.
    Global readers do not see uncommitted writes until commit().
    """

    def __init__(self, cache: LFUCache):
        self.cache = cache
        self._local_puts: Dict[str, Tuple[Any, float]] = {}  # key -> (value, absolute_expiry)
        self._local_deletes: Set[str] = set()

    async def get(self, key: str) -> Optional[Any]:
        """Read with local buffer priority (Read-Your-Own-Writes)."""
        if key in self._local_deletes:
            return None
        if key in self._local_puts:
            val, _ = self._local_puts[key]
            return val
        # Fall back to global cache (handles lazy eviction & lock)
        return await self.cache.get(key)

    async def put(self, key: str, value: Any, ttl_seconds: float) -> None:
        """Buffer write locally. Does not affect global state until commit."""
        current_time = time.monotonic()
        self._local_puts[key] = (value, current_time + ttl_seconds)
        if key in self._local_deletes:
            self._local_deletes.remove(key)

    async def delete(self, key: str) -> None:
        """Buffer deletion locally."""
        self._local_deletes.add(key)
        self._local_puts.pop(key, None)

    async def commit(self) -> None:
        """Atomically apply all buffered changes to the global cache."""
        await self.cache.apply_transaction(self)

    def rollback(self) -> None:
        """Discard all pending local changes without mutating global state."""
        self._local_puts.clear()
        self._local_deletes.clear()


# ------------------------------------------------------------------ #
#                           TEST SUITE                                 #
# ------------------------------------------------------------------ #

async def main():
    print("🧪 Starting LFU Cache Test Suite...\n")

    # a) O(1) LFU Eviction Order
    print("[a] Testing O(1) LFU Eviction Order...")
    cache = LFUCache(capacity=3)
    await cache.put("A", 1, ttl_seconds=60)
    await cache.put("B", 2, ttl_seconds=60)
    await cache.put("C", 3, ttl_seconds=60)
    
    # Access A twice to increase its frequency
    await cache.get("A")
    await cache.get("A")
    
    # Insert D. Should evict B or C (both freq=1). LFU policy guarantees one of them is gone.
    await cache.put("D", 4, ttl_seconds=60)
    
    val_b = await cache.get("B")
    val_c = await cache.get("C")
    assert val_d := await cache.get("D"), "D should exist"
    assert val_a := await cache.get("A"), "A should exist (highest freq)"
    assert val_b is None or val_c is None, f"LFU eviction failed: B={val_b}, C={val_c}"
    print(f"   ✅ LFU Eviction verified. Evicted key had lower frequency than A & D.")

    # b) Lazy TTL vs Background Async Sweep
    print("\n[b] Testing Dual-Layer TTL Eviction...")
    cache2 = LFUCache(capacity=10)
    
    # Lazy Eviction Test
    await cache2.put("lazy_key", "val", ttl_seconds=0.2)
    assert await cache2.get("lazy_key") == "val"
    await asyncio.sleep(0.3)
    assert await cache2.get("lazy_key") is None, "Lazy eviction failed"
    
    # Background Eviction Test
    await cache2.start_evictor()
    await cache2.put("bg_key", "val", ttl_seconds=0.15)
    await asyncio.sleep(0.3)  # Wait past TTL without accessing key
    assert await cache2.get("bg_key") is None, "Background async sweep failed"
    await cache2.stop_evictor()
    print("   ✅ Lazy & Background TTL eviction verified.")

    # c) Transaction Commit Visibility vs Rollback
    print("\n[c] Testing Transaction Isolation & Rollback...")
    cache3 = LFUCache(capacity=10)
    
    tx1 = cache3.begin_transaction() if hasattr(cache3, 'begin_transaction') else None
    class TxWrapper:
        def __init__(self, c): self.c = c
        def begin(self): return Transaction(self.c)
    tw = TxWrapper(cache3)

    # Commit visibility
    tx_put = tw.begin()
    await tx_put.put("committed", 100, ttl_seconds=60)
    assert await cache3.get("committed") is None, "Uncommitted write should be invisible"
    await tx_put.commit()
    assert await cache3.get("committed") == 100, "Committed write should be visible globally"

    # Rollback state restoration
    tx_roll = tw.begin()
    await tx_roll.put("rolled_back", 200, ttl_seconds=60)
    await tx_roll.rollback()
    assert await cache3.get("rolled_back") is None, "Rolled back write should not persist"
    
    # Read-Your-Own-Writes inside transaction
    tx_ryo = tw.begin()
    await tx_ryo.put("local", 999, ttl_seconds=60)
    assert await tx_ryo.get("local") == 999, "Transaction should see its own writes"
    print("   ✅ Transaction commit visibility & rollback verified.")

    # d) Stress Test: 50 Concurrent Tasks
    print("\n[d] Running Stress Test (50 concurrent async tasks)...")
    cache4 = LFUCache(capacity=100)
    await cache4.start_evictor()
    
    errors = []
    results = {"gets": 0, "puts": 0}

    async def worker(task_id: int):
        try:
            for i in range(20):
                key = f"stress_{task_id}_{i}"
                op_type = i % 3
                
                if op_type == 0:
                    await cache4.put(key, f"v_{task_id}_{i}", ttl_seconds=1.0)
                    results["puts"] += 1
                elif op_type == 1:
                    val = await cache4.get(key)
                    if val is not None:
                        results["gets"] += 1
                else:
                    tx = Transaction(cache4)
                    await tx.put(f"tx_{task_id}_{i}", "tx_val", ttl_seconds=0.5)
                    assert await tx.get(f"tx_{task_id}_{i}") == "tx_val"
                    await tx.commit()
        except Exception as e:
            errors.append((task_id, str(e)))

    tasks = [asyncio.create_task(worker(i)) for i in range(50)]
    await asyncio.gather(*tasks)
    
    assert len(errors) == 0, f"Stress test failed with errors: {errors}"
    print(f"   ✅ Stress test passed. Processed {results['puts']} puts & {results['gets']} gets across 50 tasks without race conditions.")
    await cache4.stop_evictor()

    print("\n🎉 All tests passed successfully!")


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