"""
Pluggable session storage for AgentContext.

These stores are framework-agnostic and usable anywhere:
notebooks, scripts, tests, and the WebUI all share the same backends.

    store = InMemorySessionStore()          # volatile, for tests / notebooks
    store = FileSessionStore("./sessions")  # one JSON file per session
    store = CachedFileSessionStore("./sessions")  # cache + file, recommended default
"""

from __future__ import annotations

import json
import logging
from abc import ABC, abstractmethod
from collections.abc import Iterable
from datetime import datetime
from pathlib import Path

from agentbyte.context import AgentContext

logger: logging.Logger = logging.getLogger(__name__)


class SessionStore(ABC):
    """Abstract persistence layer for AgentContext sessions."""

    @abstractmethod
    async def get(self, session_id: str) -> AgentContext | None:
        """Return a stored session context or ``None``."""

    async def save(self, session_id: str, context: AgentContext) -> None:
        """Stamp updated_at and persist the session context."""
        context.updated_at = datetime.now()
        await self._save(session_id, context)

    @abstractmethod
    async def _save(self, session_id: str, context: AgentContext) -> None:
        """Persist the full session context."""

    @abstractmethod
    async def list(self) -> Iterable[tuple[str, AgentContext]]:
        """Return all stored (session_id, context) pairs."""

    @abstractmethod
    async def delete(self, session_id: str) -> bool:
        """Delete a session if present. Returns True if it existed."""

    @abstractmethod
    async def clear_all(self) -> int:
        """Delete all stored sessions. Returns the count removed."""

    async def list_by_user(self, user_id: str) -> Iterable[tuple[str, AgentContext]]:
        """Return all (session_id, context) pairs owned by user_id.

        O(n) for file-backed stores. Override in DB-backed stores to use an index.
        """
        return [(sid, ctx) for sid, ctx in await self.list() if ctx.user_id == user_id]


class InMemorySessionStore(SessionStore):
    """Volatile in-memory session store. Lost on process restart.

    Use for: unit tests, short-lived notebook experiments.
    """

    def __init__(self) -> None:
        self._store: dict[str, AgentContext] = {}

    async def get(self, session_id: str) -> AgentContext | None:
        context = self._store.get(session_id)
        return context.model_copy(deep=True) if context is not None else None

    async def _save(self, session_id: str, context: AgentContext) -> None:
        self._store[session_id] = context.model_copy(deep=True)

    async def list(self) -> Iterable[tuple[str, AgentContext]]:
        return [
            (session_id, context.model_copy(deep=True))
            for session_id, context in self._store.items()
        ]

    async def delete(self, session_id: str) -> bool:
        return self._store.pop(session_id, None) is not None

    async def clear_all(self) -> int:
        count = len(self._store)
        self._store.clear()
        return count


class FileSessionStore(SessionStore):
    """Persistent file-based session store. Each session is one JSON file.

    Use for: durable sessions where the process may restart.
    """

    def __init__(self, storage_dir: str = ".agentbyte_sessions") -> None:
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)

    def _path(self, session_id: str) -> Path:
        return self.storage_dir / f"{session_id}.json"

    async def get(self, session_id: str) -> AgentContext | None:
        path = self._path(session_id)
        if not path.exists():
            return None
        try:
            with open(path, encoding="utf-8") as f:
                data = json.load(f)
            return AgentContext.model_validate(data)
        except Exception:
            logger.exception("Error loading session %s", session_id)
            return None

    async def _save(self, session_id: str, context: AgentContext) -> None:
        path = self._path(session_id)
        try:
            with open(path, "w", encoding="utf-8") as f:
                json.dump(context.model_dump(mode="json"), f, indent=2)
        except Exception:
            logger.exception("Error saving session %s", session_id)

    async def list(self) -> Iterable[tuple[str, AgentContext]]:
        result: list[tuple[str, AgentContext]] = []
        for path in self.storage_dir.glob("*.json"):
            context = await self.get(path.stem)
            if context is not None:
                result.append((path.stem, context))
        return result

    async def delete(self, session_id: str) -> bool:
        path = self._path(session_id)
        if path.exists():
            try:
                path.unlink()
                return True
            except Exception:
                logger.exception("Error deleting session %s", session_id)
        return False

    async def clear_all(self) -> int:
        count = 0
        for path in self.storage_dir.glob("*.json"):
            try:
                path.unlink()
                count += 1
            except Exception:
                logger.exception("Error deleting %s", path)
        return count


class CachedFileSessionStore(SessionStore):
    """In-memory cache with file persistence. Fast reads, survives restarts.

    Use for: production and long-running notebooks where you want both
    speed and durability.
    """

    def __init__(self, storage_dir: str = ".agentbyte_sessions") -> None:
        self._cache: dict[str, AgentContext] = {}
        self._file_store = FileSessionStore(storage_dir)

    async def get(self, session_id: str) -> AgentContext | None:
        if session_id in self._cache:
            return self._cache[session_id].model_copy(deep=True)
        context = await self._file_store.get(session_id)
        if context is not None:
            self._cache[session_id] = context
        return context.model_copy(deep=True) if context is not None else None

    async def _save(self, session_id: str, context: AgentContext) -> None:
        self._cache[session_id] = context.model_copy(deep=True)
        await self._file_store._save(session_id, context)

    async def list(self) -> Iterable[tuple[str, AgentContext]]:
        return await self._file_store.list()

    async def delete(self, session_id: str) -> bool:
        self._cache.pop(session_id, None)
        return await self._file_store.delete(session_id)

    async def clear_all(self) -> int:
        self._cache.clear()
        return await self._file_store.clear_all()


__all__ = [
    "CachedFileSessionStore",
    "FileSessionStore",
    "InMemorySessionStore",
    "SessionStore",
]
