import asyncio
import time
import weakref
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, List, Set, Tuple
from collections import OrderedDict
import threading
from concurrent.futures import ThreadPoolExecutor


@dataclass
class CacheEntry:
    """Represents a cache entry with value, frequency, and TTL."""
    value: Any
    freq: int = 0
    ttl: Optional[float] = None  # Unix timestamp when it expires
    created_at: float = field(default_factory=time.time)
    accessed_at: float = field(default_factory=time.time)


@dataclass
class Transaction:
    """Represents a transaction context for atomic operations."""
    cache: 'LFUCache'
    _writes: Dict[str, CacheEntry] = field(default_factory=dict)
    _deletes: Set[str] = field(default_factory=set)
    _committed: bool = False
    _rolled_back: bool = False

    def get(self, key: str) -> Optional[Any]:
        """Get value from transaction or global cache."""
        if self._rolled_back:
            raise RuntimeError("Transaction already rolled back")
        if key in self._deletes:
            return None
        if key in self._writes:
            return self._writes[key].value
        return self.cache.get(key)

    def put(self, key: str, value: Any, ttl_seconds: Optional[float] = None) -> None:
        """Put value in transaction."""
        if self._rolled_back:
            raise RuntimeError("Transaction already rolled back")
        self._writes[key] = CacheEntry(
            value=value,
            freq=0,
            ttl=(time.time() + ttl_seconds) if ttl_seconds is not None else None
        )
        self._deletes.discard(key)

    def delete(self, key: str) -> None:
        """Delete key in transaction."""
        if self._rolled_back:
            raise RuntimeError("Transaction already rolled back")
        self._deletes.add(key)
        self._writes.pop(key, None)

    async def commit(self) -> None:
        """Commit transaction to global cache."""
        if self._committed or self._rolled_back:
            return
        async with self.cache._lock:
            for key, entry in self._writes.items():
                self.cache._put_internal(key, entry)
            for key in self._deletes:
                self.cache._delete_internal(key)
        self._committed = True

    async def rollback(self) -> None:
        """Rollback transaction."""
        if self._rolled_back:
            return
        self._rolled_back = True
        self._writes.clear()
        self._deletes.clear()


class LFUCache:
    """In-Memory Concurrent LFU Cache with Async TTL Eviction and Atomic Transactions."""
    
    def __init__(self, capacity: int = 128):
        self.capacity = capacity
        self._cache: Dict[str, CacheEntry] = {}
        self._freq_buckets: Dict[int, OrderedDict[str, None]] = {}  # freq -> ordered keys
        self._key_to_freq: Dict[str, int] = {}  # key -> freq
        self._lock = asyncio.Lock()
        self._evictor_task: Optional[asyncio.Task] = None
        self._evictor_running = False
        self._executor = ThreadPoolExecutor(max_workers=1)
        
    def begin_transaction(self) -> Transaction:
        """Begin a new transaction."""
        return Transaction(self)
    
    def _get_freq_bucket(self, freq: int) -> OrderedDict[str, None]:
        """Get or create frequency bucket."""
        if freq not in self._freq_buckets:
            self._freq_buckets[freq] = OrderedDict()
        return self._freq_buckets[freq]
    
    def _update_freq(self, key: str) -> None:
        """Update frequency of a key."""
        if key not in self._key_to_freq:
            return
        old_freq = self._key_to_freq[key]
        new_freq = old_freq + 1
        self._key_to_freq[key] = new_freq
        
        # Remove from old bucket
        old_bucket = self._get_freq_bucket(old_freq)
        if key in old_bucket:
            del old_bucket[key]
        
        # Add to new bucket
        new_bucket = self._get_freq_bucket(new_freq)
        new_bucket[key] = None
        
    def _delete_internal(self, key: str) -> None:
        """Delete key from internal structures."""
        if key in self._cache:
            entry = self._cache[key]
            freq = self._key_to_freq.pop(key, 0)
            bucket = self._get_freq_bucket(freq)
            if key in bucket:
                del bucket[key]
            del self._cache[key]
    
    def _put_internal(self, key: str, entry: CacheEntry) -> None:
        """Internal put operation."""
        # Update frequency
        if key in self._cache:
            old_entry = self._cache[key]
            freq = self._key_to_freq.pop(key, 0)
            bucket = self._get_freq_bucket(freq)
            if key in bucket:
                del bucket[key]
        else:
            # New entry
            self._key_to_freq[key] = 0
        
        # Update cache and frequency tracking
        self._cache[key] = entry
        self._update_freq(key)
        
        # Evict if needed
        if len(self._cache) > self.capacity:
            self._evict_lfu()
    
    def _evict_lfu(self) -> None:
        """Evict least frequently used item."""
        # Find the lowest frequency bucket with items
        min_freq = min(self._freq_buckets.keys()) if self._freq_buckets else 0
        bucket = self._get_freq_bucket(min_freq)
        
        if bucket:
            # Remove oldest item in the lowest frequency bucket
            key = next(iter(bucket))
            self._delete_internal(key)
    
    def _is_expired(self, entry: CacheEntry) -> bool:
        """Check if an entry is expired."""
        return entry.ttl is not None and time.time() > entry.ttl
    
    def _cleanup_expired(self) -> None:
        """Cleanup expired entries."""
        expired_keys = []
        for key, entry in self._cache.items():
            if self._is_expired(entry):
                expired_keys.append(key)
        
        for key in expired_keys:
            self._delete_internal(key)
    
    async def get(self, key: str) -> Optional[Any]:
        """Get value by key with O(1) time complexity."""
        async with self._lock:
            # Cleanup expired entries
            self._cleanup_expired()
            
            if key not in self._cache:
                return None
            
            entry = self._cache[key]
            
            # Check if expired
            if self._is_expired(entry):
                self._delete_internal(key)
                return None
            
            # Update access time and frequency
            entry.accessed_at = time.time()
            self._update_freq(key)
            
            return entry.value
    
    async def put(self, key: str, value: Any, ttl_seconds: Optional[float] = None) -> None:
        """Put key-value with TTL in cache."""
        async with self._lock:
            # Cleanup expired entries
            self._cleanup_expired()
            
            entry = CacheEntry(
                value=value,
                freq=0,
                ttl=(time.time() + ttl_seconds) if ttl_seconds is not None else None
            )
            
            self._put_internal(key, entry)
    
    async def delete(self, key: str) -> None:
        """Delete a key from cache."""
        async with self._lock:
            self._delete_internal(key)
    
    def start_evictor(self) -> None:
        """Start the background evictor task."""
        if self._evictor_running:
            return
        self._evictor_running = True
        self._evictor_task = asyncio.create_task(self._evictor_loop())
    
    def stop_evictor(self) -> None:
        """Stop the background evictor task."""
        self._evictor_running = False
        if self._evictor_task:
            self._evictor_task.cancel()
    
    async def _evictor_loop(self) -> None:
        """Background task to periodically evict expired entries."""
        while self._evictor_running:
            try:
                await asyncio.sleep(1.0)
                # Run cleanup in thread pool to avoid blocking
                loop = asyncio.get_event_loop()
                await loop.run_in_executor(self._executor, self._cleanup_expired)
            except asyncio.CancelledError:
                break
            except Exception:
                pass  # Ignore errors in evictor loop


async def main():
    """Run unit tests for the LFUCache implementation."""
    
    print("Starting LFU Cache Tests...")
    
    # Test 1: O(1) LFU eviction order
    print("\n1. Testing O(1) LFU eviction...")
    cache = LFUCache(capacity=3)
    
    # Put 4 items, should evict the least frequently used
    await cache.put("a", "value_a")
    await cache.put("b", "value_b")
    await cache.put("c", "value_c")
    
    # Access "a" to increase its frequency
    await cache.get("a")
    await cache.put("d", "value_d")  # Should evict "b" (least frequent)
    
    assert await cache.get("a") == "value_a", "Should still have 'a'"
    assert await cache.get("b") is None, "Should have evicted 'b'"
    assert await cache.get("c") == "value_c", "Should still have 'c'"
    assert await cache.get("d") == "value_d", "Should have 'd'"
    
    print("✓ LFU eviction works correctly")
    
    # Test 2: Lazy TTL vs. Background Async Sweep
    print("\n2. Testing TTL eviction...")
    cache = LFUCache(capacity=10)
    
    # Put with short TTL
    await cache.put("short", "value", ttl_seconds=0.1)
    await cache.put("long", "value2", ttl_seconds=1.0)
    
    # Should have both entries
    assert await cache.get("short") == "value"
    assert await cache.get("long") == "value2"
    
    # Wait for short TTL to expire
    await asyncio.sleep(0.15)
    
    # Short should be gone, long should still be there
    assert await cache.get("short") is None, "Short TTL should have expired"
    assert await cache.get("long") == "value2", "Long TTL should still be valid"
    
    print("✓ TTL eviction works correctly")
    
    # Test 3: Transaction commit visibility vs rollback
    print("\n3. Testing transaction isolation...")
    cache = LFUCache(capacity=10)
    
    # Start transaction
    tx = cache.begin_transaction()
    
    # Write to transaction
    tx.put("tx_key", "tx_value")
    assert tx.get("tx_key") == "tx_value", "Transaction should see its own writes"
    
    # Global cache should not see transaction changes
    assert await cache.get("tx_key") is None, "Global cache should not see uncommitted changes"
    
    # Commit transaction
    await tx.commit()
    
    # Global cache should now see changes
    assert await cache.get("tx_key") == "tx_value", "Global cache should see committed changes"
    
    # Start another transaction
    tx2 = cache.begin_transaction()
    tx2.put("tx_key2", "tx_value2")
    
    # Rollback
    await tx2.rollback()
    
    # Should not see rolled back changes
    assert await cache.get("tx_key2") is None, "Should not see rolled back changes"
    
    print("✓ Transaction isolation works correctly")
    
    # Test 4: Stress test with concurrent tasks
    print("\n4. Testing concurrent access...")
    cache = LFUCache(capacity=50)
    
    async def worker(worker_id: int, num_ops: int):
        for i in range(num_ops):
            key = f"worker_{worker_id}_key_{i}"
            await cache.put(key, f"value_{worker_id}_{i}")
            value = await cache.get(key)
            assert value is not None, f"Worker {worker_id} failed to get value for key {key}"
    
    # Run 50 concurrent tasks
    tasks = [worker(i, 10) for i in range(50)]
    await asyncio.gather(*tasks)
    
    # Verify all keys are present
    for i in range(50):
        for j in range(10):
            key = f"worker_{i}_key_{j}"
            value = await cache.get(key)
            assert value is not None, f"Key {key} should exist"
    
    print("✓ Concurrent access works correctly")
    
    # Cleanup
    cache.stop_evictor()
    print("\nAll tests passed!")


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