"""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.
"""

from __future__ import annotations

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

from core.types import ThreadState

if TYPE_CHECKING:
    from persistence.store import StudioStore

logger = logging.getLogger(__name__)


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


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

    Handles two types of state:
    1. Thread state - Conversation history (agent_id + thread_id)
    2. Shared memory - Key-value store with namespaces
    """

    # -------------------------------------------------------------------------
    # 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."""
        ...


# =============================================================================
# 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)


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("_")}


# =============================================================================
# STORE WITH STUDIO STORE BACKEND
# =============================================================================


class PersistentStateStore(StateStore):
    """State store backed by StudioStore for production use.

    Uses the existing persistence layer for threads.
    Shared memory delegates to a single FileStateStore instance.

    Conversion between ``ThreadState`` (core domain) and ``ThreadRecord``
    (persistence boundary) is done through explicit field mapping in
    ``_to_record()`` / ``_to_state()``, not opaque dict round-trips.
    """

    def __init__(self, store: StudioStore):
        self._store = store
        self._lock = asyncio.Lock()
        # Single shared file store for shared memory (not per-call)
        self._file_store = FileStateStore()

    # -------------------------------------------------------------------------
    # ThreadState ↔ ThreadRecord conversions
    # -------------------------------------------------------------------------

    @staticmethod
    def _to_record(state: ThreadState, *, key: str) -> object:
        """Convert ThreadState → ThreadRecord for the persistence boundary."""
        from persistence.store import ThreadRecord  # noqa: PLC0415 (circular)

        return ThreadRecord(
            id=key,
            agent_id=state.agent_id,
            thread_id=state.thread_id,
            messages=[m.model_dump(mode="json") for m in state.messages],
            working_memory=state.working_memory,
            metadata=state.metadata,
            last_active=state.last_active,
            turn_count=state.turn_count,
            total_tokens_used=state.total_tokens_used,
            schema_version=state.schema_version,
        )

    @staticmethod
    def _to_state(record: object, agent_id: str, thread_id: str) -> ThreadState:
        """Convert ThreadRecord → ThreadState from the persistence boundary."""
        from core.types.messages import ConversationMessage

        raw = record.model_dump(mode="json")  # type: ignore[union-attr]
        messages = [
            ConversationMessage.model_validate(m)
            for m in (raw.get("messages") or [])
        ]
        return ThreadState(
            agent_id=raw.get("agent_id") or agent_id,
            thread_id=raw.get("thread_id") or thread_id,
            messages=messages,
            working_memory=raw.get("working_memory") or {},
            metadata=raw.get("metadata") or {},
            last_active=raw.get("last_active"),
            turn_count=raw.get("turn_count", 0),
            total_tokens_used=raw.get("total_tokens_used", 0),
            schema_version=raw.get("schema_version", 1),
        )

    # -------------------------------------------------------------------------
    # Thread State (via StudioStore)
    # -------------------------------------------------------------------------

    async def load_thread(self, agent_id: str, thread_id: str) -> ThreadState:
        async with self._lock:
            key = f"{agent_id}:{thread_id}"
            record = await asyncio.to_thread(self._store.get_thread, key)
            if not record:
                return ThreadState(agent_id=agent_id, thread_id=thread_id)
            return self._to_state(record, agent_id, thread_id)

    async def save_thread(self, state: ThreadState) -> None:
        async with self._lock:
            state.last_active = datetime.now(UTC).isoformat()
            key = f"{state.agent_id}:{state.thread_id}"
            record = self._to_record(state, key=key)
            await asyncio.to_thread(self._store.save_thread, key, record)

    async def delete_thread(self, agent_id: str, thread_id: str) -> bool:
        async with self._lock:
            key = f"{agent_id}:{thread_id}"
            return await asyncio.to_thread(self._store.delete_thread, key)

    async def list_threads(
        self,
        *,
        agent_id: str | None = None,
        limit: int = 100,
    ) -> list[ThreadState]:
        async with self._lock:
            threads = await asyncio.to_thread(self._store.list_threads, limit=limit)
            states: list[ThreadState] = []
            for t in threads:
                try:
                    raw = t.model_dump(mode="json")
                    state = self._to_state(
                        t,
                        agent_id=raw.get("agent_id", ""),
                        thread_id=raw.get("thread_id", ""),
                    )
                    if agent_id is None or state.agent_id == agent_id:
                        states.append(state)
                except Exception:
                    continue
            return states

    # -------------------------------------------------------------------------
    # Shared Memory (delegates to FileStateStore)
    # -------------------------------------------------------------------------

    async def get_shared(self, key: str, namespace: str = "default") -> Any | None:
        return await self._file_store.get_shared(key, namespace)

    async def set_shared(self, key: str, value: Any, namespace: str = "default") -> None:
        await self._file_store.set_shared(key, value, namespace)

    async def delete_shared(self, key: str, namespace: str = "default") -> bool:
        return await self._file_store.delete_shared(key, namespace)

    async def list_shared(self, namespace: str = "default") -> dict[str, Any]:
        return await self._file_store.list_shared(namespace)


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