"""In-memory backend implementation for testing and development."""

from typing import Optional, List, Dict
from datetime import datetime

from .base import StateBackend
from ..schemas import UnifiedState, AgentState, StateMetadata


class MemoryBackend(StateBackend):
    """In-memory state storage backend (for testing/development)."""
    
    def __init__(self):
        """Initialize in-memory backend."""
        self._storage: Dict[str, UnifiedState] = {}
        self._shared_storage: Dict[str, UnifiedState] = {}
        self._agents: Dict[str, set] = {}  # session_id -> set of agent_ids
    
    def write_state(self, state: UnifiedState) -> bool:
        """Write state to memory."""
        key = state.get_key()
        self._storage[key] = state
        
        # Track agents in session
        if state.session_id not in self._agents:
            self._agents[state.session_id] = set()
        self._agents[state.session_id].add(state.agent_id)
        
        return True
    
    def read_state(self, session_id: str, agent_id: str) -> Optional[UnifiedState]:
        """Read state from memory."""
        key = f"uass:session:{session_id}:agent:{agent_id}:state"
        return self._storage.get(key)
    
    def delete_state(self, session_id: str, agent_id: str) -> bool:
        """Delete state from memory."""
        key = f"uass:session:{session_id}:agent:{agent_id}:state"
        if key in self._storage:
            del self._storage[key]
            if session_id in self._agents:
                self._agents[session_id].discard(agent_id)
            return True
        return False
    
    def list_agents(self, session_id: str) -> List[str]:
        """List all agent IDs in a session."""
        return list(self._agents.get(session_id, set()))
    
    def read_shared_state(self, session_id: str) -> Optional[UnifiedState]:
        """Read shared state for a session."""
        return self._shared_storage.get(session_id)
    
    def write_shared_state(self, session_id: str, state: UnifiedState) -> bool:
        """Write shared state for a session."""
        self._shared_storage[session_id] = state
        return True
    
    def clear(self):
        """Clear all stored state (for testing)."""
        self._storage.clear()
        self._shared_storage.clear()
        self._agents.clear()

