"""State Store - Unified state management.

Handles:
- Thread state (conversation history per agent+thread)
- Shared memory (key-value scratch space, shared across agents and turns)

All async methods use asyncio.Lock and asyncio.to_thread for file I/O
to avoid blocking the event loop.
"""

import asyncio
import json
import logging
from abc import ABC, abstractmethod
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from core.types import ThreadState

logger = logging.getLogger(__name__)


# =============================================================================
# STATE STORE PROTOCOL
# =============================================================================


class StateStore(ABC):
    """Unified state storage.

    Handles three types of state:
    1. Thread state - Conversation history (agent_id + thread_id)
    2. Shared memory - Key-value store with namespaces
    3. DAG checkpoints - Resumable workflow snapshots
    """

    # -------------------------------------------------------------------------
    # Thread State
    # -------------------------------------------------------------------------

    @abstractmethod
    async def load_thread(self, agent_id: str, thread_id: str) -> ThreadState:
        """Load thread state. Returns empty state if not found."""
        ...

    @abstractmethod
    async def save_thread(self, state: ThreadState) -> None:
        """Save thread state."""
        ...

    @abstractmethod
    async def delete_thread(self, agent_id: str, thread_id: str) -> bool:
        """Delete thread state."""
        ...

    @abstractmethod
    async def list_threads(
        self,
        *,
        agent_id: str | None = None,
        limit: int = 100,
    ) -> list[ThreadState]:
        """List thread states."""
        ...

    # -------------------------------------------------------------------------
    # Shared Memory (key-value with namespaces)
    # -------------------------------------------------------------------------

    @abstractmethod
    async def get_shared(self, key: str, namespace: str = "default") -> Any | None:
        """Get a value from shared memory."""
        ...

    @abstractmethod
    async def set_shared(self, key: str, value: Any, namespace: str = "default") -> None:
        """Set a value in shared memory."""
        ...

    @abstractmethod
    async def delete_shared(self, key: str, namespace: str = "default") -> bool:
        """Delete a value from shared memory."""
        ...

    @abstractmethod
    async def list_shared(self, namespace: str = "default") -> dict[str, Any]:
        """List all key-values in a shared namespace."""
        ...

    # -------------------------------------------------------------------------
    # DAG Checkpoints
    # -------------------------------------------------------------------------

    async def save_checkpoint(self, checkpoint: Any) -> None:  # noqa: B027
        """Save a DAG checkpoint for resume-from-failure.

        Default implementation is a no-op — subclasses opt in.
        ``FileStateStore`` provides a concrete implementation.
        """

    async def load_checkpoint(self, execution_id: str) -> Any | None:  # noqa: ARG002
        """Load a DAG checkpoint by *execution_id*.

        Returns ``None`` if no checkpoint exists.  Default is no-op.
        """
        return None

    async def delete_checkpoint(self, execution_id: str) -> bool:  # noqa: ARG002
        """Delete a DAG checkpoint.  Default is no-op."""
        return False


# =============================================================================
# FILE-BASED IMPLEMENTATION
# =============================================================================


def _ensure_dirs(root: Path) -> None:
    """Create storage directories (sync, called once at init)."""
    (root / "threads").mkdir(parents=True, exist_ok=True)
    (root / "shared").mkdir(parents=True, exist_ok=True)
    (root / "checkpoints").mkdir(parents=True, exist_ok=True)


class FileStateStore(StateStore):
    """File-based state store.

    Storage layout:
        .albus_store/
            threads/
                {agent_id}_{thread_id}.json
            shared/
                {namespace}.json

    All file I/O is offloaded via asyncio.to_thread to avoid blocking
    the event loop. An asyncio.Lock serializes access within each store.
    """

    def __init__(self, root: str = ".albus_store"):
        self._root = Path(root)
        self._lock = asyncio.Lock()
        _ensure_dirs(self._root)

    # -------------------------------------------------------------------------
    # Thread State
    # -------------------------------------------------------------------------

    def _thread_path(self, agent_id: str, thread_id: str) -> Path:
        safe_id = f"{agent_id}_{thread_id}".replace("/", "_").replace("\\", "_")
        return self._root / "threads" / f"{safe_id}.json"

    def _load_thread_sync(self, agent_id: str, thread_id: str) -> ThreadState:
        path = self._thread_path(agent_id, thread_id)
        if not path.exists():
            return ThreadState(agent_id=agent_id, thread_id=thread_id)
        try:
            data = json.loads(path.read_text())
            return ThreadState.model_validate(data)
        except Exception as e:
            logger.warning("Failed to load thread state: %s", e)
            return ThreadState(agent_id=agent_id, thread_id=thread_id)

    def _save_thread_sync(self, state: ThreadState) -> None:
        state.last_active = datetime.now(UTC).isoformat()
        path = self._thread_path(state.agent_id, state.thread_id)
        path.write_text(json.dumps(state.model_dump(mode="json"), indent=2))

    async def load_thread(self, agent_id: str, thread_id: str) -> ThreadState:
        async with self._lock:
            return await asyncio.to_thread(self._load_thread_sync, agent_id, thread_id)

    async def save_thread(self, state: ThreadState) -> None:
        async with self._lock:
            await asyncio.to_thread(self._save_thread_sync, state)

    async def delete_thread(self, agent_id: str, thread_id: str) -> bool:
        async with self._lock:
            path = self._thread_path(agent_id, thread_id)
            if not path.exists():
                return False
            await asyncio.to_thread(path.unlink)
            return True

    async def list_threads(
        self,
        *,
        agent_id: str | None = None,
        limit: int = 100,
    ) -> list[ThreadState]:
        def _list_sync() -> list[ThreadState]:
            threads_dir = self._root / "threads"
            if not threads_dir.exists():
                return []
            states: list[ThreadState] = []
            for path in threads_dir.glob("*.json"):
                try:
                    data = json.loads(path.read_text())
                    state = ThreadState.model_validate(data)
                    if agent_id is None or state.agent_id == agent_id:
                        states.append(state)
                except Exception:
                    continue
            states.sort(key=lambda s: s.last_active or "", reverse=True)
            return states[:limit]

        async with self._lock:
            return await asyncio.to_thread(_list_sync)

    # -------------------------------------------------------------------------
    # Shared Memory
    # -------------------------------------------------------------------------

    def _shared_path(self, namespace: str) -> Path:
        safe_ns = namespace.replace("/", "_").replace("\\", "_")
        return self._root / "shared" / f"{safe_ns}.json"

    def _load_shared_ns_sync(self, namespace: str) -> dict[str, Any]:
        path = self._shared_path(namespace)
        if not path.exists():
            return {}
        try:
            return json.loads(path.read_text())
        except Exception:
            return {}

    def _save_shared_ns_sync(self, namespace: str, data: dict[str, Any]) -> None:
        path = self._shared_path(namespace)
        path.write_text(json.dumps(data, indent=2))

    async def get_shared(self, key: str, namespace: str = "default") -> Any | None:
        async with self._lock:
            data = await asyncio.to_thread(self._load_shared_ns_sync, namespace)
            return data.get(key)

    async def set_shared(self, key: str, value: Any, namespace: str = "default") -> None:
        async with self._lock:
            data = await asyncio.to_thread(self._load_shared_ns_sync, namespace)
            data[key] = value
            data["_updated_at"] = datetime.now(UTC).isoformat()
            await asyncio.to_thread(self._save_shared_ns_sync, namespace, data)

    async def delete_shared(self, key: str, namespace: str = "default") -> bool:
        async with self._lock:
            data = await asyncio.to_thread(self._load_shared_ns_sync, namespace)
            if key not in data:
                return False
            del data[key]
            data["_updated_at"] = datetime.now(UTC).isoformat()
            await asyncio.to_thread(self._save_shared_ns_sync, namespace, data)
            return True

    async def list_shared(self, namespace: str = "default") -> dict[str, Any]:
        async with self._lock:
            data = await asyncio.to_thread(self._load_shared_ns_sync, namespace)
            return {k: v for k, v in data.items() if not k.startswith("_")}

    # -------------------------------------------------------------------------
    # DAG Checkpoints
    # -------------------------------------------------------------------------

    def _checkpoint_path(self, execution_id: str) -> Path:
        safe_id = execution_id.replace("/", "_").replace("\\", "_")
        return self._root / "checkpoints" / f"{safe_id}.json"

    async def save_checkpoint(self, checkpoint: Any) -> None:
        """Persist a DAG checkpoint to disk."""
        from core.pathways.checkpoint import DAGCheckpoint

        if not isinstance(checkpoint, DAGCheckpoint):
            return
        path = self._checkpoint_path(checkpoint.execution_id)

        async with self._lock:
            await asyncio.to_thread(
                path.write_text,
                json.dumps(checkpoint.model_dump(), indent=2),
            )

    async def load_checkpoint(self, execution_id: str) -> Any | None:
        """Load a DAG checkpoint from disk."""
        from core.pathways.checkpoint import DAGCheckpoint

        path = self._checkpoint_path(execution_id)
        if not path.exists():
            return None

        async with self._lock:
            text = await asyncio.to_thread(path.read_text)
            return DAGCheckpoint.model_validate_json(text)

    async def delete_checkpoint(self, execution_id: str) -> bool:
        """Delete a DAG checkpoint from disk."""
        path = self._checkpoint_path(execution_id)
        if not path.exists():
            return False
        async with self._lock:
            await asyncio.to_thread(path.unlink)
            return True


__all__ = [
    "FileStateStore",
    "StateStore",
]
