# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/1-Session.ipynb.

# %% auto 0
__all__ = ['Message', 'ChatSession', 'SessionManager', 'InMemorySessionManager']

# %% ../../nbs/1-Session.ipynb 1
from typing import List, Dict, Optional, Any, Protocol
from dataclasses import dataclass, field
from abc import ABC, abstractmethod
import uuid

# %% ../../nbs/1-Session.ipynb 3
@dataclass
class Message:
    """Represents a single chat message."""

    role: str
    content: str
    id: Optional[str] = None
    pending: bool = False
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        if self.id is None:
            self.id = str(uuid.uuid4())

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary format for rendering."""
        result = {
            "role": self.role,
            "content": self.content,
            "id": self.id,
        }
        if self.pending:
            result["pending"] = self.pending
        result.update(self.metadata)
        return result

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Message":
        """Create Message from dictionary."""
        metadata = {
            k: v
            for k, v in data.items()
            if k not in ("role", "content", "id", "pending")
        }
        return cls(
            role=data["role"],
            content=data.get("content", ""),
            id=data.get("id"),
            pending=data.get("pending", False),
            metadata=metadata,
        )

# %% ../../nbs/1-Session.ipynb 4
class ChatSession:
    """Manages a single chat conversation."""

    def __init__(
        self, session_id: str, initial_messages: Optional[List[Message]] = None
    ):
        self.session_id = session_id
        self._messages: List[Message] = initial_messages or []

    def add_message(
        self, role: str, content: str, pending: bool = False, **metadata
    ) -> Message:
        """Add a new message to the session."""
        msg = Message(role=role, content=content, pending=pending, metadata=metadata)
        self._messages.append(msg)
        return msg

    def update_message(self, msg_id: str, **updates) -> Optional[Message]:
        """Update an existing message by ID."""
        for msg in self._messages:
            if msg.id == msg_id:
                for key, value in updates.items():
                    if key == "content":
                        msg.content = value
                    elif key == "pending":
                        msg.pending = value
                    elif key == "role":
                        msg.role = value
                    else:
                        msg.metadata[key] = value
                return msg
        return None

    def remove_message(self, msg_id: str) -> bool:
        """Remove a message by ID."""
        original_len = len(self._messages)
        self._messages = [m for m in self._messages if m.id != msg_id]
        return len(self._messages) < original_len

    def get_messages(self) -> List[Message]:
        """Get all messages in the session."""
        return self._messages.copy()

    def get_message_dicts(self) -> List[Dict[str, Any]]:
        """Get all messages as dictionaries."""
        return [msg.to_dict() for msg in self._messages]

    def clear(self):
        """Clear all messages."""
        self._messages.clear()

    def __len__(self) -> int:
        return len(self._messages)

# %% ../../nbs/1-Session.ipynb 5
class SessionManager(ABC):
    """Abstract base class for managing multiple chat sessions."""

    @abstractmethod
    def create_session(
        self,
        session_id: Optional[str] = None,
        initial_messages: Optional[List[Message]] = None,
    ) -> ChatSession:
        """Create a new chat session."""
        pass

    @abstractmethod
    def get_session(self, session_id: str) -> Optional[ChatSession]:
        """Retrieve a session by ID."""
        pass

    @abstractmethod
    def delete_session(self, session_id: str) -> bool:
        """Delete a session by ID."""
        pass

    @abstractmethod
    def list_sessions(self) -> List[str]:
        """List all session IDs."""
        pass

# %% ../../nbs/1-Session.ipynb 6
class InMemorySessionManager(SessionManager):
    """In-memory implementation of SessionManager."""

    def __init__(self):
        self._sessions: Dict[str, ChatSession] = {}

    def create_session(
        self,
        session_id: Optional[str] = None,
        initial_messages: Optional[List[Message]] = None,
    ) -> ChatSession:
        """Create a new chat session."""
        if session_id is None:
            session_id = str(uuid.uuid4())

        session = ChatSession(session_id, initial_messages)
        self._sessions[session_id] = session
        return session

    def get_session(self, session_id: str) -> Optional[ChatSession]:
        """Retrieve a session by ID."""
        return self._sessions.get(session_id)

    def delete_session(self, session_id: str) -> bool:
        """Delete a session by ID."""
        if session_id in self._sessions:
            del self._sessions[session_id]
            return True
        return False

    def list_sessions(self) -> List[str]:
        """List all session IDs."""
        return list(self._sessions.keys())

    def __len__(self) -> int:
        return len(self._sessions)
