"""Session storage implementations.
SessionStore protocol and FileSessionStore implementation.
"""
from __future__ import annotations
import json
import os
from dataclasses import asdict, dataclass, replace
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Protocol
import trio
if TYPE_CHECKING:
    pass
from .dtypes import (
    AgentSession,
    Endpoint,
    EnvironmentConfig,
    Message,
    SessionStatus,
)
def _load_first_user_message(messages_file: Path) -> list[Message]:
    """Load first user message from messages.jsonl for preview."""
    if not messages_file.exists():
        return []
    with open(messages_file) as f:
        for raw_line in f:
            stripped = raw_line.strip()
            if not stripped:
                continue
            msg_data = json.loads(stripped)
            if msg_data.get("role") == "user":
                return [Message(**msg_data)]
    return []
def _load_session_preview(session_dir: Path) -> AgentSession | None:
    """Load session metadata and first user message for listing.
    Returns None if session cannot be loaded (invalid/corrupt).
    Pure function: directory in, session out.
    """
    session_file = session_dir / "session.json"
    if not session_file.exists():
        return None
    try:
        with open(session_file) as f:
            data = json.load(f)
        messages = _load_first_user_message(session_dir / "messages.jsonl")
        return AgentSession(
            session_id=data.get("session_id", session_dir.name),
            endpoint=Endpoint(provider="", model=data.get("model", "")),
            status=SessionStatus(data.get("status", "active")),
            tags=data.get("tags", {}),
            messages=messages,
            created_at=data.get("created_at", datetime.now().isoformat()),
            updated_at=data.get("updated_at", datetime.now().isoformat()),
        )
    except Exception:
        return None
def generate_session_id() -> str:
    """Generate a unique session ID.
    Format: timestamp_random (e.g., "20241205_143052_a1b2c3")
    """
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    random_suffix = os.urandom(3).hex()
    return f"{timestamp}_{random_suffix}"
class SessionStore(Protocol):
    """Storage backend for AgentSessions.
    Implementations should be frozen dataclasses (just config, no mutable state).
    This allows SessionStore to be serializable and passed around freely.
    """
    # Core CRUD
    async def create(
        self,
        endpoint: Endpoint,
        environment: EnvironmentConfig,
        parent_id: str | None = None,
        branch_point: int | None = None,
        tags: dict[str, str] | None = None,
    ) -> AgentSession:
        """Create new session, return AgentSession with generated session_id."""
        ...
    async def get(self, session_id: str) -> tuple[AgentSession | None, str | None]:
        """Load session by ID. Returns (session, None) or (None, error)."""
        ...
    async def update(
        self,
        session_id: str,
        status: SessionStatus | None = None,
        environment_state: dict | None = None,
        reward: float | dict[str, float] | None = None,
        tags: dict[str, str] | None = None,
        endpoint: Endpoint | None = None,
    ) -> tuple[None, str | None]:
        """Update session metadata. Returns (None, None) or (None, error)."""
        ...
    # Streaming append
    async def append_message(self, session_id: str, message: Message) -> None:
        """Append message to trajectory (streaming, append-only)."""
        ...
    # Queries
    async def list_sessions(
        self,
        filter_tags: dict[str, str] | None = None,
        status: SessionStatus | None = None,
        limit: int = 100,
    ) -> list[AgentSession]:
        """List sessions, optionally filtered by tags and status."""
        ...
    async def list_children(self, parent_id: str) -> list[AgentSession]:
        """List child sessions (branches/resumes)."""
        ...
    async def get_latest(
        self,
        status: SessionStatus | None = None,
    ) -> tuple[AgentSession | None, str | None]:
        """Get the most recent session, optionally filtered by status."""
        ...
    # Cleanup
    async def delete(self, session_id: str) -> tuple[None, str | None]:
        """Delete session and associated data."""
        ...
@dataclass(frozen=True)
class FileSessionStore:
    """File-based implementation of SessionStore.
    Frozen dataclass - just holds the base_dir path. Serializable.
    Methods are essentially pure functions that take base_dir as implicit arg.
    Layout:
        ~/.rollouts/sessions/
            <session_id>/
                session.json     # metadata: endpoint, environment, tags, status, parent_id, etc.
                messages.jsonl   # trajectory (append-only)
                environment.json # serialized env state (written at checkpoints)
    """
    base_dir: Path = Path.home() / ".rollouts" / "sessions"
    def _ensure_base_dir(self) -> None:
        """Lazy directory creation - called by methods that need it.
        Avoids side effects in __post_init__ which would break frozen dataclass semantics.
        Idempotent due to exist_ok=True.
        """
        self.base_dir.mkdir(parents=True, exist_ok=True)
    def _session_dir(self, session_id: str) -> Path:
        """Get the directory for a session."""
        return self.base_dir / session_id
    async def _write_json(self, path: Path, data: dict) -> None:
        """Write JSON to file atomically."""
        content = json.dumps(data, indent=2)
        # Write to temp file then rename for atomicity
        temp_path = path.with_suffix(".tmp")
        async with await trio.open_file(temp_path, "w") as f:
            await f.write(content)
        temp_path.rename(path)
    async def _read_json(self, path: Path) -> dict:
        """Read JSON from file."""
        try:
            async with await trio.open_file(path, "r") as f:
                content = await f.read()
            return json.loads(content)
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON in {path}: {e}") from e
        except Exception as e:
            raise IOError(f"Failed to read {path}: {e}") from e
    async def create(
        self,
        endpoint: Endpoint,
        environment: EnvironmentConfig,
        parent_id: str | None = None,
        branch_point: int | None = None,
        tags: dict[str, str] | None = None,
    ) -> AgentSession:
        """Create new session, return AgentSession with generated session_id."""
        self._ensure_base_dir()
        session_id = generate_session_id()
        session_dir = self._session_dir(session_id)
        session_dir.mkdir()
        now = datetime.now().isoformat()
        session = AgentSession(
            session_id=session_id,
            parent_id=parent_id,
            branch_point=branch_point,
            endpoint=endpoint,
            environment=environment,
            messages=[],
            status=SessionStatus.PENDING,
            tags=tags or {},
            created_at=now,
            updated_at=now,
        )
        # Write session.json
        await self._write_json(session_dir / "session.json", session.to_dict())
        (session_dir / "messages.jsonl").touch()
        return session
    async def save(self, session: AgentSession) -> None:
        """Save a complete session (including messages).
        Used for saving transformed sessions (compact, summarize).
        """
        self._ensure_base_dir()
        session_dir = self._session_dir(session.session_id)
        session_dir.mkdir(exist_ok=True)
        # Write session.json
        await self._write_json(session_dir / "session.json", session.to_dict())
        # Write messages.jsonl
        messages_file = session_dir / "messages.jsonl"
        async with await trio.open_file(messages_file, "w") as f:
            for msg in session.messages:
                await f.write(json.dumps(asdict(msg)) + "\n")
    def get_config_sync(self, session_id: str) -> tuple[dict | None, str | None]:
        """Sync load of session config (just session.json, not messages).
        Used during CLI arg parsing before async context is available.
        Returns (config_dict, None) or (None, error).
        """
        session_dir = self._session_dir(session_id)
        if not session_dir.exists():
            return None, f"Session not found: {session_id}"
        session_file = session_dir / "session.json"
        if not session_file.exists():
            return None, f"Session config not found: {session_id}"
        return json.loads(session_file.read_text()), None
    def get_latest_id_sync(self) -> str | None:
        """Sync get the most recent session ID.
        Used during CLI arg parsing before async context is available.
        """
        self._ensure_base_dir()
        if not self.base_dir.exists():
            return None
        session_dirs = sorted(self.base_dir.iterdir(), reverse=True)
        for session_dir in session_dirs:
            if session_dir.is_dir() and (session_dir / "session.json").exists():
                return session_dir.name
        return None
    def list_sync(self, limit: int = 20) -> list[AgentSession]:
        """Sync list recent sessions.
        Used for --list-sessions before async context is available.
        Returns sessions sorted by ID (newest first), without messages.
        """
        self._ensure_base_dir()
        if not self.base_dir.exists():
            return []
        sessions: list[AgentSession] = []
        session_dirs = sorted(self.base_dir.iterdir(), reverse=True)
        for session_dir in session_dirs[:limit]:
            if not session_dir.is_dir():
                continue
            session = _load_session_preview(session_dir)
            if session is not None:
                sessions.append(session)
        return sessions
    async def get(self, session_id: str) -> tuple[AgentSession | None, str | None]:
        """Load session by ID. Returns (session, None) or (None, error)."""
        session_dir = self._session_dir(session_id)
        if not session_dir.exists():
            return None, f"Session not found: {session_id}"
        try:
            session_data = await self._read_json(session_dir / "session.json")
        except Exception as e:
            return None, f"Failed to read session.json: {e}"
        messages: list[Message] = []
        messages_file = session_dir / "messages.jsonl"
        if messages_file.exists():
            try:
                async with await trio.open_file(messages_file, "r") as f:
                    async for raw_line in f:
                        stripped = raw_line.strip()
                        if stripped:
                            # Fail fast if message can't be parsed - corrupted data is a bug
                            messages.append(Message.from_json(stripped))
            except Exception as e:
                return None, f"Failed to read messages.jsonl: {e}"
        try:
            return AgentSession.from_dict(session_data, messages), None
        except Exception as e:
            return None, f"Failed to create AgentSession: {e}"
    async def update(
        self,
        session_id: str,
        status: SessionStatus | None = None,
        environment_state: dict | None = None,
        reward: float | dict[str, float] | None = None,
        tags: dict[str, str] | None = None,
        endpoint: Endpoint | None = None,
    ) -> tuple[None, str | None]:
        """Update session metadata. Returns (None, None) or (None, error)."""
        session_dir = self._session_dir(session_id)
        if not session_dir.exists():
            return None, f"Session not found: {session_id}"
        session_data = await self._read_json(session_dir / "session.json")

        if status is not None:
            session_data["status"] = status.value
        if environment_state is not None:
            session_data["environment_state"] = environment_state
        if reward is not None:
            session_data["reward"] = reward
        if tags is not None:
            session_data["tags"] = tags
        if endpoint is not None:
            session_data["endpoint"] = endpoint.to_dict(exclude_secrets=True)
        session_data["updated_at"] = datetime.now().isoformat()
        # Write back
        await self._write_json(session_dir / "session.json", session_data)
        return None, None
    async def append_message(self, session_id: str, message: Message) -> None:
        """Append message to trajectory (streaming, append-only)."""
        session_dir = self._session_dir(session_id)
        messages_file = session_dir / "messages.jsonl"
        if message.timestamp is None:
            message = replace(message, timestamp=datetime.now().isoformat())
        # Append-only (streaming safe)
        async with await trio.open_file(messages_file, "a") as f:
            await f.write(message.to_json() + "\n")
    async def list_sessions(
        self,
        filter_tags: dict[str, str] | None = None,
        status: SessionStatus | None = None,
        limit: int = 100,
    ) -> list[AgentSession]:
        """List sessions, optionally filtered by tags and status."""
        self._ensure_base_dir()
        sessions: list[AgentSession] = []
        # Iterate through session directories
        if not self.base_dir.exists():
            return sessions
        session_dirs = sorted(self.base_dir.iterdir(), reverse=True)  # newest first
        for session_dir in session_dirs:
            if not session_dir.is_dir():
                continue
            session_file = session_dir / "session.json"
            if not session_file.exists():
                continue
            async with await trio.open_file(session_file, "r") as f:
                content = await f.read()
            session_data = json.loads(content)
            if status is not None and session_data.get("status") != status.value:
                continue
            if filter_tags:
                session_tags = session_data.get("tags", {})
                if not all(session_tags.get(k) == v for k, v in filter_tags.items()):
                    continue
            # Count messages without loading them
            messages_file = session_dir / "messages.jsonl"
            message_count = 0
            if messages_file.exists():
                async with await trio.open_file(messages_file, "r") as f:
                    async for line in f:
                        if line.strip():
                            message_count += 1
            session = AgentSession.from_dict(session_data, messages=[])
            session = replace(session, message_count=message_count)
            sessions.append(session)
            if len(sessions) >= limit:
                break
        return sessions
    async def list_children(self, parent_id: str) -> list[AgentSession]:
        """List child sessions (branches/resumes)."""
        self._ensure_base_dir()
        children: list[AgentSession] = []
        if not self.base_dir.exists():
            return children
        for session_dir in self.base_dir.iterdir():
            if not session_dir.is_dir():
                continue
            session_file = session_dir / "session.json"
            if not session_file.exists():
                continue
            async with await trio.open_file(session_file, "r") as f:
                content = await f.read()
            session_data = json.loads(content)
            if session_data.get("parent_id") == parent_id:
                session = AgentSession.from_dict(session_data, messages=[])
                children.append(session)
        return children
    async def get_latest(
        self,
        status: SessionStatus | None = None,
    ) -> tuple[AgentSession | None, str | None]:
        """Get the most recent session, optionally filtered by status."""
        sessions = await self.list_sessions(status=status, limit=1)
        if not sessions:
            return None, "No sessions found"
        return sessions[0], None
    async def delete(self, session_id: str) -> tuple[None, str | None]:
        """Delete session and associated data."""
        import shutil
        session_dir = self._session_dir(session_id)
        if not session_dir.exists():
            return None, f"Session not found: {session_id}"
        # Remove entire directory
        shutil.rmtree(session_dir)
        return None, None
def log_crash(
    error: Exception,
    provider: str,
    model: str,
    *,
    session_id: str | None = None,
    messages: list | None = None,
) -> Path:
    """Log crash info to ~/.rollouts/crashes/ with optional message dump.
    Returns path to the crash file for reference in error messages.
    """
    import traceback
    crashes_dir = Path.home() / ".rollouts" / "crashes"
    crashes_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    random_suffix = os.urandom(3).hex()
    crash_file = crashes_dir / f"{timestamp}_{random_suffix}.txt"
    crash_info: dict = {
        "timestamp": datetime.now().isoformat(),
        "provider": provider,
        "model": model,
        "session_id": session_id,
        "error_type": type(error).__name__,
        "error_message": str(error),
        "traceback": traceback.format_exc(),
    }
    if messages is not None:
        crash_info["messages"] = messages
    crash_file.write_text(json.dumps(crash_info, indent=2, default=str))
    return crash_file
