from typing import Any, List, Dict
import asyncio

from xulcan.memory.state.base import BaseStateStore
from xulcan.core.primitives import MachineID

class MemoryStateStore(BaseStateStore):
    """
    In-Memory implementation of the State Store.
    Uses Python dictionaries to hold state. Extremely fast, but volatile.
    Perfect for Standalone mode or single-node deployments.
    """
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Estructura: _storage[run_id][key] = value
        self._storage: Dict[MachineID, Dict[str, Any]] = {}
        # Un lock global simple para evitar condiciones de carrera raras en asyncio
        self._lock = asyncio.Lock()

    async def _set_impl(self, run_id: MachineID, key: str, value: Any) -> None:
        async with self._lock:
            if run_id not in self._storage:
                self._storage[run_id] = {}
            self._storage[run_id][key] = value

    async def _get_impl(self, run_id: MachineID, key: str) -> Any:
        async with self._lock:
            run_data = self._storage.get(run_id, {})
            return run_data.get(key, None)

    async def _exists_impl(self, run_id: MachineID, key: str) -> bool:
        async with self._lock:
            return run_id in self._storage and key in self._storage[run_id]

    async def _keys_impl(self, run_id: MachineID) -> List[str]:
        async with self._lock:
            if run_id not in self._storage:
                return []
            return list(self._storage[run_id].keys())

    async def _delete_impl(self, run_id: MachineID, key: str) -> None:
        async with self._lock:
            if run_id in self._storage and key in self._storage[run_id]:
                del self._storage[run_id][key]

    async def _clear_impl(self, run_id: MachineID) -> None:
        async with self._lock:
            if run_id in self._storage:
                del self._storage[run_id]