"""
Working Memory System for Toxo.

Implements working memory for temporary storage and active processing.
"""

import asyncio
from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from datetime import datetime
from collections import deque

from ..utils.logger import get_logger


@dataclass
class WorkingMemoryItem:
    """A working memory item."""
    id: str
    content: str
    priority: float = 0.5
    timestamp: datetime = None
    ttl: float = 300.0  # Time to live in seconds
    
    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = datetime.now()
        elif isinstance(self.timestamp, str):
            self.timestamp = datetime.fromisoformat(self.timestamp)
    
    @property
    def is_expired(self) -> bool:
        """Check if item has expired."""
        age = (datetime.now() - self.timestamp).total_seconds()
        return age > self.ttl


class WorkingMemorySystem:
    """Working memory system for temporary storage."""
    
    def __init__(self, capacity: int = 100):
        self.logger = get_logger("toxo.memory.working")
        self.capacity = capacity
        self.items: Dict[str, WorkingMemoryItem] = {}
        self.access_queue: deque = deque(maxlen=capacity)
        self.logger.info(f"Working memory system initialized with capacity {capacity}")
    
    async def store(self, data: Dict[str, Any]) -> str:
        """Store item in working memory."""
        content = data.get("content", "")
        priority = data.get("priority", 0.5)
        ttl = data.get("ttl", 300.0)
        
        item_id = f"working_{len(self.items)}_{datetime.now().timestamp()}"
        
        item = WorkingMemoryItem(
            id=item_id,
            content=content,
            priority=priority,
            ttl=ttl
        )
        
        # Clean expired items
        await self._cleanup_expired()
        
        # Make room if needed
        if len(self.items) >= self.capacity:
            await self._evict_item()
        
        self.items[item_id] = item
        self.access_queue.append(item_id)
        
        self.logger.debug(f"Stored working memory item: {item_id}")
        return item_id
    
    async def retrieve(self, query: str) -> List[Dict[str, Any]]:
        """Retrieve items from working memory."""
        await self._cleanup_expired()
        
        results = []
        for item in self.items.values():
            if not item.is_expired and query.lower() in item.content.lower():
                results.append({
                    "id": item.id,
                    "content": item.content,
                    "priority": item.priority,
                    "timestamp": item.timestamp.isoformat(),
                    "ttl": item.ttl
                })
        
        # Sort by priority and recency
        results.sort(key=lambda x: (x["priority"], x["timestamp"]), reverse=True)
        return results
    
    async def search(self, query: str) -> List[Dict[str, Any]]:
        """Search working memory."""
        return await self.retrieve(query)
    
    async def analyze(self) -> Dict[str, Any]:
        """Analyze working memory."""
        await self._cleanup_expired()
        return {
            "total_items": len(self.items),
            "capacity": self.capacity,
            "utilization": len(self.items) / self.capacity,
            "avg_priority": sum(item.priority for item in self.items.values()) / max(1, len(self.items))
        }
    
    async def optimize(self) -> Dict[str, Any]:
        """Optimize working memory."""
        await self._cleanup_expired()
        self.logger.info("Optimizing working memory")
        return {"status": "optimized", "items_count": len(self.items)}
    
    async def consolidate(self) -> Dict[str, Any]:
        """Consolidate working memory."""
        await self._cleanup_expired()
        self.logger.info("Consolidating working memory")
        return {"status": "consolidated", "items_count": len(self.items)}
    
    async def _cleanup_expired(self):
        """Remove expired items."""
        expired_ids = [item_id for item_id, item in self.items.items() if item.is_expired]
        for item_id in expired_ids:
            del self.items[item_id]
            if item_id in self.access_queue:
                self.access_queue.remove(item_id)
    
    async def _evict_item(self):
        """Evict least recently used item."""
        if self.access_queue:
            lru_id = self.access_queue.popleft()
            if lru_id in self.items:
                del self.items[lru_id]
    
    def get_stats(self) -> Dict[str, Any]:
        """Get memory statistics."""
        return {
            "type": "working",
            "items": len(self.items),
            "capacity": self.capacity,
            "utilization": len(self.items) / self.capacity
        } 