"""Synchronous MongoDB storage implementation for Upsonic agent framework."""
from __future__ import annotations

import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

if TYPE_CHECKING:
    from pymongo import MongoClient
    from pymongo.collection import Collection
    from pymongo.database import Database
    from upsonic.session.base import Session, SessionType
    from upsonic.culture.cultural_knowledge import CulturalKnowledge

try:
    from pymongo import MongoClient, ReturnDocument, ReplaceOne
    from pymongo.collection import Collection
    from pymongo.database import Database

    PYMONGO_AVAILABLE = True
except ImportError:
    PYMONGO_AVAILABLE = False
    MongoClient = None  # type: ignore
    Collection = None  # type: ignore
    Database = None  # type: ignore
    ReturnDocument = None  # type: ignore
    ReplaceOne = None  # type: ignore

from upsonic.storage.base import Storage
from upsonic.storage.mongo.utils import (
    apply_pagination,
    apply_sorting,
    create_collection_indexes,
    remove_mongo_id,
)
from upsonic.storage.schemas import UserMemory
from upsonic.storage.utils import deserialize_session, deserialize_session_json_fields
from upsonic.utils.logging_config import get_logger


_logger = get_logger("upsonic.storage.mongo")


class MongoStorage(Storage):
    """
    Synchronous MongoDB storage implementation for Upsonic agent framework.
    
    This storage backend provides persistence for agent sessions and user memory data
    using PyMongo's synchronous MongoClient.
    
    Example:
        ```python
        storage = MongoStorage(
            db_url="mongodb://localhost:27017",
            db_name="upsonic_db",
        )
        storage._create_all_tables()
        
        # Upsert a session
        session = AgentSession(session_id="abc123", ...)
        result = storage.upsert_session(session)
        
        # Get a session
        session = storage.get_session(session_id="abc123")
        ```
    """

    def __init__(
        self,
        db_client: Optional["MongoClient"] = None,
        db_name: Optional[str] = None,
        db_url: Optional[str] = None,
        session_collection: Optional[str] = None,
        user_memory_collection: Optional[str] = None,
        id: Optional[str] = None,
    ) -> None:
        """
        Initialize the synchronous MongoDB storage.
        
        Args:
            db_client: The MongoDB client to use. If not provided, a client
                      will be created from db_url.
            db_name: The name of the database to use.
            db_url: The database URL to connect to.
            session_collection: Name of the collection to store sessions.
            user_memory_collection: Name of the collection to store user memories.
            id: Unique identifier for this storage instance.
        
        Raises:
            ValueError: If neither db_url nor db_client is provided.
            ImportError: If pymongo is not installed.
        """
        if not PYMONGO_AVAILABLE:
            from upsonic.utils.printing import import_error
            import_error(
                package_name="pymongo",
                install_command='pip install "upsonic[mongo-storage]"',
                feature_name="MongoDB storage provider"
            )
        
        super().__init__(
            session_table=session_collection,
            user_memory_table=user_memory_collection,
            id=id,
        )

        # Store configuration
        self.db_url: Optional[str] = db_url
        self.db_name: str = db_name if db_name is not None else "upsonic"

        if db_client is None and db_url is None:
            raise ValueError("One of db_url or db_client must be provided")

        # Initialize or create client
        if db_client is not None:
            self._client: MongoClient = db_client
        else:
            self._client = MongoClient(db_url)

        # Database and collection caches
        self._database: Optional[Database] = None
        self._session_collection: Optional[Collection] = None
        self._user_memory_collection: Optional[Collection] = None
        self._cultural_knowledge_collection: Optional[Collection] = None
        self._sessions_initialized: bool = False
        self._user_memories_initialized: bool = False
        self._cultural_knowledge_initialized: bool = False

    # ======================== Client Management ========================

    @property
    def db_client(self) -> "MongoClient":
        """Get the MongoDB client."""
        return self._client

    @property
    def database(self) -> "Database":
        """Get the MongoDB database."""
        if self._database is None:
            self._database = self._client[self.db_name]
        return self._database

    def close(self) -> None:
        """
        Close the MongoDB client connection.
        
        Should be called during application shutdown to properly release
        all database connections.
        """
        if self._client is not None:
            self._client.close()
            self._client = None  # type: ignore
            self._database = None
            self._session_collection = None
            self._user_memory_collection = None
            self._cultural_knowledge_collection = None

    # ======================== Table/Collection Management ========================

    def table_exists(self, table_name: str) -> bool:
        """
        Check if a collection with the given name exists in the MongoDB database.
        
        Args:
            table_name: Name of the collection to check
        
        Returns:
            True if the collection exists in the database, False otherwise
        """
        collection_names = self.database.list_collection_names()
        return table_name in collection_names

    def _create_all_tables(self) -> None:
        """Create all configured MongoDB collections if they don't exist."""
        collections_to_create = [
            ("sessions", self.session_table_name),
            ("user_memories", self.user_memory_table_name),
            ("cultural_knowledge", self.cultural_knowledge_table_name),
        ]

        for collection_type, collection_name in collections_to_create:
            if collection_name and not self.table_exists(collection_name):
                self._get_collection(
                    collection_type, create_collection_if_not_found=True
                )

    def _get_collection(
        self, collection_type: str, create_collection_if_not_found: bool = True
    ) -> Optional["Collection"]:
        """
        Get or create a collection based on collection type.
        
        Args:
            collection_type: The type of collection to get ('sessions' or 'user_memories')
            create_collection_if_not_found: Whether to create the collection if it doesn't exist
        
        Returns:
            The collection object
        
        Raises:
            ValueError: If collection type is unknown or collection name not provided
        """
        if collection_type == "sessions":
            if self._session_collection is None:
                if self.session_table_name is None:
                    raise ValueError(
                        "Session collection was not provided on initialization"
                    )
                self._session_collection = self._get_or_create_collection(
                    collection_name=self.session_table_name,
                    collection_type="sessions",
                    create_collection_if_not_found=create_collection_if_not_found,
                )
            return self._session_collection

        if collection_type == "user_memories":
            if self._user_memory_collection is None:
                if self.user_memory_table_name is None:
                    raise ValueError(
                        "User memory collection was not provided on initialization"
                    )
                self._user_memory_collection = self._get_or_create_collection(
                    collection_name=self.user_memory_table_name,
                    collection_type="user_memories",
                    create_collection_if_not_found=create_collection_if_not_found,
                )
            return self._user_memory_collection

        if collection_type == "cultural_knowledge":
            if self._cultural_knowledge_collection is None:
                if self.cultural_knowledge_table_name is None:
                    raise ValueError(
                        "Cultural knowledge collection was not provided on initialization"
                    )
                self._cultural_knowledge_collection = self._get_or_create_collection(
                    collection_name=self.cultural_knowledge_table_name,
                    collection_type="cultural_knowledge",
                    create_collection_if_not_found=create_collection_if_not_found,
                )
            return self._cultural_knowledge_collection

        raise ValueError(f"Unknown collection type: {collection_type}")

    def _get_or_create_collection(
        self,
        collection_name: str,
        collection_type: str,
        create_collection_if_not_found: bool = True,
    ) -> Optional["Collection"]:
        """
        Get or create a collection with proper indexes.
        
        Args:
            collection_name: The name of the collection to get or create
            collection_type: The type of collection ('sessions' or 'user_memories')
            create_collection_if_not_found: Whether to create the collection if it doesn't exist
        
        Returns:
            The collection object
        """
        try:
            # Check if database is available (may be None after disconnect)
            if self.database is None:
                if not create_collection_if_not_found:
                    return None
                raise ValueError("Database connection is not available")
            
            collection = self.database[collection_name]

            # Check if already initialized
            init_flag = f"_{collection_type}_initialized"
            if not getattr(self, init_flag, False):
                if not create_collection_if_not_found:
                    return None
                # Create indexes synchronously
                create_collection_indexes(collection, collection_type)
                setattr(self, init_flag, True)
                _logger.debug(f"Initialized collection '{collection_name}'")
            else:
                _logger.debug(f"Collection '{collection_name}' already initialized")

            return collection

        except Exception as e:
            _logger.error(f"Error getting collection {collection_name}: {e}")
            raise

    def _get_session_type_value(self, session: "Session") -> str:
        """
        Extract session type value string from session object.
        
        Args:
            session: Session object
        
        Returns:
            Session type string value
        """
        session_type_value = "agent"
        if hasattr(session, "session_type"):
            st = session.session_type
            if hasattr(st, "value"):
                session_type_value = st.value
            elif isinstance(st, str):
                session_type_value = st
        return session_type_value

    # ======================== Session Methods ========================

    def upsert_session(
        self,
        session: "Session",
        deserialize: bool = True,
    ) -> Optional[Union["Session", Dict[str, Any]]]:
        """
        Insert or update a session in the database.
        
        Args:
            session: The session to upsert. Must have session_id set.
            deserialize: If True, return deserialized Session object.
                        If False, return raw dictionary.
        
        Returns:
            The upserted session (deserialized or dict based on flag).
        
        Raises:
            ValueError: If session_id is not provided.
        """
        from upsonic.session.agent import AgentSession
        from upsonic.session.base import SessionType

        if not hasattr(session, "session_id") or session.session_id is None:
            raise ValueError("Session must have session_id set for upsert")

        try:
            collection = self._get_collection(
                "sessions", create_collection_if_not_found=True
            )
            if collection is None:
                return None

            session_dict = (
                session.to_dict(serialize_flag=True) if hasattr(session, "to_dict") else dict(session)
            )
            current_time = int(time.time())

            if isinstance(session, AgentSession):
                record = {
                    "session_id": session_dict.get("session_id"),
                    "session_type": SessionType.AGENT.value,
                    "agent_id": session_dict.get("agent_id"),
                    "team_id": session_dict.get("team_id"),
                    "workflow_id": session_dict.get("workflow_id"),
                    "user_id": session_dict.get("user_id"),
                    "session_data": session_dict.get("session_data"),
                    "agent_data": session_dict.get("agent_data"),
                    "metadata": session_dict.get("metadata"),
                    "runs": session_dict.get("runs"),
                    "messages": session_dict.get("messages"),
                    "summary": session_dict.get("summary"),
                    "usage": session_dict.get("usage"),
                    "created_at": session_dict.get("created_at") or current_time,
                    "updated_at": current_time,
                }
            else:
                # Fallback for other session types
                session_type_value = self._get_session_type_value(session)
                record = {
                    "session_id": session_dict.get("session_id"),
                    "session_type": session_type_value,
                    "agent_id": session_dict.get("agent_id"),
                    "team_id": session_dict.get("team_id"),
                    "workflow_id": session_dict.get("workflow_id"),
                    "user_id": session_dict.get("user_id"),
                    "session_data": session_dict.get("session_data"),
                    "agent_data": session_dict.get("agent_data"),
                    "team_data": session_dict.get("team_data"),
                    "workflow_data": session_dict.get("workflow_data"),
                    "metadata": session_dict.get("metadata"),
                    "runs": session_dict.get("runs"),
                    "messages": session_dict.get("messages"),
                    "summary": session_dict.get("summary"),
                    "usage": session_dict.get("usage"),
                    "created_at": session_dict.get("created_at") or current_time,
                    "updated_at": current_time,
                }

            result = collection.find_one_and_replace(
                filter={"session_id": record["session_id"]},
                replacement=record,
                upsert=True,
                return_document=ReturnDocument.AFTER,
            )

            if result is None:
                return None

            session_raw = deserialize_session_json_fields(result)

            if not deserialize:
                return session_raw

            if isinstance(session, AgentSession):
                return AgentSession.from_dict(session_raw, deserialize_flag=True)
            return deserialize_session(session_raw)

        except Exception as e:
            _logger.error(f"Error upserting session: {e}")
            raise

    def upsert_sessions(
        self,
        sessions: List["Session"],
        deserialize: bool = True,
    ) -> List[Union["Session", Dict[str, Any]]]:
        """
        Bulk insert or update multiple sessions.
        
        Args:
            sessions: List of sessions to upsert. Each must have session_id set.
            deserialize: If True, return deserialized Session objects.
                        If False, return raw dictionaries.
        
        Returns:
            List of upserted sessions.
        
        Raises:
            ValueError: If any session is missing session_id.
        """
        if not sessions:
            return []

        # Validate all sessions have session_id
        for session in sessions:
            if not hasattr(session, "session_id") or session.session_id is None:
                raise ValueError("All sessions must have session_id set for upsert")

        try:
            collection = self._get_collection(
                "sessions", create_collection_if_not_found=True
            )
            if collection is None:
                _logger.info(
                    "Sessions collection not available, falling back to individual upserts"
                )
                return [
                    result
                    for session in sessions
                    if session is not None
                    for result in [self.upsert_session(session, deserialize=deserialize)]
                    if result is not None
                ]

            operations = []
            current_time = int(time.time())

            for session in sessions:
                if session is None:
                    continue

                session_dict = (
                    session.to_dict() if hasattr(session, "to_dict") else dict(session)
                )
                session_type_value = self._get_session_type_value(session)

                record = {
                    "session_id": session_dict.get("session_id"),
                    "session_type": session_type_value,
                    "agent_id": session_dict.get("agent_id"),
                    "team_id": session_dict.get("team_id"),
                    "workflow_id": session_dict.get("workflow_id"),
                    "user_id": session_dict.get("user_id"),
                    "session_data": session_dict.get("session_data"),
                    "agent_data": session_dict.get("agent_data"),
                    "team_data": session_dict.get("team_data"),
                    "workflow_data": session_dict.get("workflow_data"),
                    "metadata": session_dict.get("metadata"),
                    "runs": session_dict.get("runs"),
                    "messages": session_dict.get("messages"),
                    "summary": session_dict.get("summary"),
                    "usage": session_dict.get("usage"),
                    "created_at": session_dict.get("created_at") or current_time,
                    "updated_at": current_time,
                }

                operations.append(
                    ReplaceOne(
                        filter={"session_id": record["session_id"]},
                        replacement=record,
                        upsert=True,
                    )
                )

            results: List[Union["Session", Dict[str, Any]]] = []

            if operations:
                # Execute bulk write
                collection.bulk_write(operations)

                # Fetch the results
                session_ids = [
                    session.session_id
                    for session in sessions
                    if session and session.session_id
                ]
                cursor = collection.find({"session_id": {"$in": session_ids}})

                for doc in cursor:
                    session_raw = deserialize_session_json_fields(doc)

                    if deserialize:
                        results.append(deserialize_session(session_raw))
                    else:
                        results.append(session_raw)

            return results

        except Exception as e:
            _logger.error(
                f"Exception during bulk session upsert, falling back to individual upserts: {e}"
            )
            return [
                result
                for session in sessions
                if session is not None
                for result in [self.upsert_session(session, deserialize=deserialize)]
                if result is not None
            ]

    def get_session(
        self,
        session_id: Optional[str] = None,
        session_type: Optional["SessionType"] = None,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        deserialize: bool = True,
    ) -> Optional[Union["Session", Dict[str, Any]]]:
        """
        Get a session from the database.
        
        Args:
            session_id: ID of the session to retrieve. If None, returns latest session.
            session_type: Type of session (AGENT, TEAM, WORKFLOW).
            user_id: Filter by user ID.
            agent_id: Filter by agent ID.
            deserialize: If True, return deserialized Session object.
                        If False, return raw dictionary.
        
        Returns:
            The session if found, None otherwise.
        """
        try:
            collection = self._get_collection(
                "sessions", create_collection_if_not_found=True
            )
            if collection is None:
                return None
        except ValueError:
            return None

        try:
            query: Dict[str, Any] = {}

            if session_id is not None:
                query["session_id"] = session_id
            if session_type is not None:
                query["session_type"] = session_type.value
            if user_id is not None:
                query["user_id"] = user_id
            if agent_id is not None:
                query["agent_id"] = agent_id

            # If no session_id provided, get latest
            if session_id is None:
                cursor = collection.find(query).sort("created_at", -1).limit(1)
                result = list(cursor)
                if not result:
                    return None
                result = result[0]
            else:
                result = collection.find_one(query)

            if result is None:
                return None

            session_raw = deserialize_session_json_fields(result)

            if not deserialize:
                return session_raw

            return deserialize_session(session_raw, session_type)

        except Exception as e:
            _logger.error(f"Error getting session: {e}")
            raise

    def get_sessions(
        self,
        session_ids: Optional[List[str]] = None,
        session_type: Optional["SessionType"] = None,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        sort_by: Optional[str] = None,
        sort_order: Optional[str] = None,
        deserialize: bool = True,
    ) -> Union[List["Session"], Tuple[List[Dict[str, Any]], int]]:
        """
        Get multiple sessions from the database.
        
        Args:
            session_ids: List of session IDs to retrieve. If None, returns all sessions.
            session_type: Filter by session type.
            user_id: Filter by user ID.
            agent_id: Filter by agent ID.
            limit: Maximum number of sessions to return.
            offset: Number of sessions to skip.
            sort_by: Field to sort by.
            sort_order: Sort order ('asc' or 'desc').
            deserialize: If True, return deserialized Session objects.
                        If False, return tuple of (raw dictionaries, total count).
        
        Returns:
            List of sessions or tuple of (list of dicts, total count).
        """
        try:
            collection = self._get_collection(
                "sessions", create_collection_if_not_found=True
            )
            if collection is None:
                return [] if deserialize else ([], 0)
        except ValueError:
            return [] if deserialize else ([], 0)

        try:
            # If empty list is provided, return empty results immediately
            if session_ids is not None and len(session_ids) == 0:
                return [] if deserialize else ([], 0)

            query: Dict[str, Any] = {}

            if session_ids is not None and len(session_ids) > 0:
                query["session_id"] = {"$in": session_ids}
            if session_type is not None:
                query["session_type"] = session_type.value
            if user_id is not None:
                query["user_id"] = user_id
            if agent_id is not None:
                query["agent_id"] = agent_id

            # Get total count
            total_count = collection.count_documents(query)

            # Build cursor
            cursor = collection.find(query)

            # Apply sorting
            sort_criteria = apply_sorting({}, sort_by, sort_order)
            if sort_criteria:
                cursor = cursor.sort(sort_criteria)

            # Apply pagination
            cursor = apply_pagination(cursor, limit, offset)

            records = list(cursor)

            if not records:
                return [] if deserialize else ([], 0)

            sessions_raw = [
                deserialize_session_json_fields(record) for record in records
            ]

            if not deserialize:
                return sessions_raw, total_count

            return [
                deserialize_session(s, session_type) for s in sessions_raw
            ]

        except Exception as e:
            _logger.error(f"Error getting sessions: {e}")
            raise

    def delete_session(self, session_id: str) -> bool:
        """
        Delete a session from the database.
        
        Args:
            session_id: ID of the session to delete.
        
        Returns:
            True if deleted successfully, False otherwise.
        
        Raises:
            ValueError: If session_id is not provided.
        """
        if not session_id:
            raise ValueError("session_id is required for delete")

        try:
            collection = self._get_collection(
                "sessions", create_collection_if_not_found=False
            )
            if collection is None:
                return False
        except ValueError:
            return False

        try:
            result = collection.delete_one({"session_id": session_id})
            deleted = result.deleted_count > 0
            if deleted:
                _logger.debug(f"Deleted session: {session_id}")
            else:
                _logger.debug(f"No session found to delete: {session_id}")
            return deleted

        except Exception as e:
            _logger.error(f"Error deleting session: {e}")
            raise

    def delete_sessions(self, session_ids: List[str]) -> int:
        """
        Delete multiple sessions from the database.
        
        Args:
            session_ids: List of session IDs to delete.
        
        Returns:
            Number of sessions deleted.
        
        Raises:
            ValueError: If session_ids is empty.
        """
        if not session_ids:
            raise ValueError("session_ids is required and cannot be empty")

        try:
            collection = self._get_collection(
                "sessions", create_collection_if_not_found=False
            )
            if collection is None:
                return 0
        except ValueError:
            return 0

        try:
            result = collection.delete_many({"session_id": {"$in": session_ids}})
            deleted_count = result.deleted_count
            _logger.debug(f"Deleted {deleted_count} sessions")
            return deleted_count

        except Exception as e:
            _logger.error(f"Error deleting sessions: {e}")
            raise

    # ======================== User Memory Methods ========================

    def upsert_user_memory(
        self,
        user_memory: UserMemory,
        deserialize: bool = True,
    ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
        """
        Insert or update user memory in the database.
        
        Args:
            user_memory: The UserMemory instance to store. Must have user_id set.
            deserialize: If True, return full record. If False, return raw dict.
        
        Returns:
            The upserted user memory record.
        
        Raises:
            ValueError: If user_id is not provided in user_memory.
        """
        if not user_memory.user_id:
            raise ValueError("user_memory must have user_id set for upsert")

        try:
            collection = self._get_collection(
                "user_memories", create_collection_if_not_found=True
            )
            if collection is None:
                return None

            current_time = int(time.time())

            record = {
                "user_id": user_memory.user_id,
                "user_memory": user_memory.user_memory,
                "agent_id": user_memory.agent_id,
                "team_id": user_memory.team_id,
                "created_at": user_memory.created_at or current_time,
                "updated_at": current_time,
            }

            result = collection.find_one_and_replace(
                filter={"user_id": user_memory.user_id},
                replacement=record,
                upsert=True,
                return_document=ReturnDocument.AFTER,
            )

            if result is None:
                return None

            memory_dict = remove_mongo_id(result)
            if deserialize:
                return UserMemory.from_dict(memory_dict)
            return memory_dict

        except Exception as e:
            _logger.error(f"Error upserting user memory: {e}")
            raise

    def upsert_user_memories(
        self,
        user_memories: List[UserMemory],
        deserialize: bool = True,
    ) -> List[Union[UserMemory, Dict[str, Any]]]:
        """
        Bulk insert or update multiple user memories.
        
        Args:
            user_memories: List of UserMemory instances. Each must have user_id set.
            deserialize: If True, return full records.
        
        Returns:
            List of upserted user memory records.
        
        Raises:
            ValueError: If any record is missing user_id.
        """
        if not user_memories:
            return []

        # Validate all have user_id
        for memory in user_memories:
            if not memory.user_id:
                raise ValueError("All user memories must have user_id set")

        try:
            collection = self._get_collection(
                "user_memories", create_collection_if_not_found=True
            )
            if collection is None:
                _logger.info(
                    "User memories collection not available, falling back to individual upserts"
                )
                return [
                    result
                    for memory in user_memories
                    if memory is not None
                    for result in [
                        self.upsert_user_memory(
                            user_memory=memory,
                            deserialize=deserialize,
                        )
                    ]
                    if result is not None
                ]

            operations = []
            current_time = int(time.time())

            for memory in user_memories:
                if memory is None:
                    continue

                record = {
                    "user_id": memory.user_id,
                    "user_memory": memory.user_memory,
                    "agent_id": memory.agent_id,
                    "team_id": memory.team_id,
                    "created_at": memory.created_at or current_time,
                    "updated_at": current_time,
                }

                operations.append(
                    ReplaceOne(
                        filter={"user_id": memory.user_id},
                        replacement=record,
                        upsert=True,
                    )
                )

            results: List[Union[UserMemory, Dict[str, Any]]] = []

            if operations:
                # Execute bulk write
                collection.bulk_write(operations)

                # Fetch the results
                user_ids = [
                    memory.user_id
                    for memory in user_memories
                    if memory and memory.user_id
                ]
                cursor = collection.find({"user_id": {"$in": user_ids}})

                for doc in cursor:
                    memory_dict = remove_mongo_id(doc)
                    if deserialize:
                        results.append(UserMemory.from_dict(memory_dict))
                    else:
                        results.append(memory_dict)

            return results

        except Exception as e:
            _logger.error(
                f"Exception during bulk user memory upsert, falling back to individual upserts: {e}"
            )
            return [
                result
                for memory in user_memories
                if memory is not None
                for result in [
                    self.upsert_user_memory(
                        user_memory=memory,
                        deserialize=deserialize,
                    )
                ]
                if result is not None
            ]

    def get_user_memory(
        self,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        team_id: Optional[str] = None,
        deserialize: bool = True,
    ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
        """
        Get user memory from the database.
        
        Args:
            user_id: User ID to retrieve. If None, returns latest memory.
            agent_id: Filter by agent ID.
            team_id: Filter by team ID.
            deserialize: If True, return UserMemory object. If False, return raw dict.
        
        Returns:
            The user memory record if found, None otherwise.
        """
        try:
            collection = self._get_collection(
                "user_memories", create_collection_if_not_found=True
            )
            if collection is None:
                return None
        except ValueError:
            return None

        try:
            query: Dict[str, Any] = {}

            if user_id is not None:
                query["user_id"] = user_id
            if agent_id is not None:
                query["agent_id"] = agent_id
            if team_id is not None:
                query["team_id"] = team_id

            # If no user_id provided, get latest
            if user_id is None:
                cursor = collection.find(query).sort("updated_at", -1).limit(1)
                result = list(cursor)
                if not result:
                    return None
                result = result[0]
            else:
                result = collection.find_one(query)

            if result is None:
                return None

            memory_dict = remove_mongo_id(result)
            if deserialize:
                return UserMemory.from_dict(memory_dict)
            return memory_dict

        except Exception as e:
            _logger.error(f"Error getting user memory: {e}")
            raise

    def get_user_memories(
        self,
        user_ids: Optional[List[str]] = None,
        agent_id: Optional[str] = None,
        team_id: Optional[str] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        deserialize: bool = True,
    ) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
        """
        Get multiple user memories from the database.
        
        Args:
            user_ids: List of user IDs to retrieve. If None, returns all.
            agent_id: Filter by agent ID.
            team_id: Filter by team ID.
            limit: Maximum number of records to return.
            offset: Number of records to skip.
            deserialize: If True, return list of UserMemory objects. 
                         If False, return tuple of (list of dicts, total count).
        
        Returns:
            List of UserMemory objects if deserialize=True, otherwise tuple of (dicts, count).
        """
        try:
            collection = self._get_collection(
                "user_memories", create_collection_if_not_found=False
            )
            if collection is None:
                if deserialize:
                    return []
                return [], 0
        except ValueError:
            if deserialize:
                return []
            return [], 0

        try:
            # If empty list is provided, return empty results immediately
            if user_ids is not None and len(user_ids) == 0:
                if deserialize:
                    return []
                return [], 0

            query: Dict[str, Any] = {}

            if user_ids is not None and len(user_ids) > 0:
                query["user_id"] = {"$in": user_ids}
            if agent_id is not None:
                query["agent_id"] = agent_id
            if team_id is not None:
                query["team_id"] = team_id

            # Get total count
            total_count = collection.count_documents(query)

            # Build cursor
            cursor = collection.find(query).sort("updated_at", -1)

            # Apply pagination
            cursor = apply_pagination(cursor, limit, offset)

            records = list(cursor)

            if not records:
                if deserialize:
                    return []
                return [], 0

            if deserialize:
                return [
                    UserMemory.from_dict(remove_mongo_id(record)) for record in records
                ]
            return [
                remove_mongo_id(record) for record in records
            ], total_count

        except Exception as e:
            _logger.error(f"Error getting user memories: {e}")
            raise

    def delete_user_memory(self, user_id: str) -> bool:
        """
        Delete user memory from the database.
        
        Args:
            user_id: User ID of the memory to delete.
        
        Returns:
            True if deleted successfully, False otherwise.
        
        Raises:
            ValueError: If user_id is not provided.
        """
        if not user_id:
            raise ValueError("user_id is required for delete")

        try:
            collection = self._get_collection(
                "user_memories", create_collection_if_not_found=False
            )
            if collection is None:
                return False
        except ValueError:
            return False

        try:
            result = collection.delete_one({"user_id": user_id})
            deleted = result.deleted_count > 0
            if deleted:
                _logger.debug(f"Deleted user memory: {user_id}")
            else:
                _logger.debug(f"No user memory found to delete: {user_id}")
            return deleted

        except Exception as e:
            _logger.error(f"Error deleting user memory: {e}")
            raise

    def delete_user_memories(self, user_ids: List[str]) -> int:
        """
        Delete multiple user memories from the database.
        
        Args:
            user_ids: List of user IDs to delete.
        
        Returns:
            Number of records deleted.
        
        Raises:
            ValueError: If user_ids is empty.
        """
        if not user_ids:
            raise ValueError("user_ids is required and cannot be empty")

        try:
            collection = self._get_collection(
                "user_memories", create_collection_if_not_found=False
            )
            if collection is None:
                return 0
        except ValueError:
            return 0

        try:
            result = collection.delete_many({"user_id": {"$in": user_ids}})
            deleted_count = result.deleted_count
            _logger.debug(f"Deleted {deleted_count} user memories")
            return deleted_count

        except Exception as e:
            _logger.error(f"Error deleting user memories: {e}")
            raise

    # ======================== Utility Methods ========================

    def clear_all(self) -> None:
        """
        Clear all data from all collections.
        
        This removes all sessions and user memories from the storage.
        Use with caution.
        """
        try:
            # Clear sessions
            try:
                collection = self._get_collection(
                    "sessions", create_collection_if_not_found=False
                )
                if collection is not None:
                    collection.delete_many({})
                    _logger.debug("Cleared all sessions")
            except ValueError:
                pass

            # Clear user memories
            try:
                collection = self._get_collection(
                    "user_memories", create_collection_if_not_found=False
                )
                if collection is not None:
                    collection.delete_many({})
                    _logger.debug("Cleared all user memories")
            except ValueError:
                pass

            _logger.info("Cleared all data from storage")

        except Exception as e:
            _logger.error(f"Error clearing all data: {e}")
            raise

    # ======================== Generic Model Methods ========================

    def _get_generic_model_collection(self, collection: str) -> "Collection":
        """Get or create a collection for generic models."""
        collection_name = f"upsonic_models_{collection}"
        return self.db[collection_name]

    def upsert_model(
        self,
        key: str,
        model: Any,
        collection: str = "generic_models",
    ) -> None:
        """Insert or update a generic Pydantic model in storage."""
        import time
        
        try:
            coll = self._get_generic_model_collection(collection)
            current_time = int(time.time())
            
            # Serialize model
            if hasattr(model, 'model_dump'):
                model_data = model.model_dump(mode='json')
            elif hasattr(model, 'dict'):
                model_data = model.dict()
            else:
                model_data = dict(model)
            
            document = {
                "_id": key,
                "collection": collection,
                "model_data": model_data,
                "updated_at": current_time,
            }
            
            # Check if exists to preserve created_at
            existing = coll.find_one({"_id": key})
            if existing:
                document["created_at"] = existing.get("created_at", current_time)
            else:
                document["created_at"] = current_time
            
            coll.replace_one({"_id": key}, document, upsert=True)
            _logger.debug(f"Upserted model with key '{key}' in collection '{collection}'")
            
        except Exception as e:
            _logger.error(f"Error upserting model: {e}")
            raise

    def get_model(
        self,
        key: str,
        model_type: Any,
        collection: str = "generic_models",
    ) -> Optional[Any]:
        """Retrieve a generic Pydantic model from storage."""
        try:
            coll = self._get_generic_model_collection(collection)
            
            document = coll.find_one({"_id": key})
            if document is None:
                return None
            
            model_data = document.get("model_data")
            if model_data is None:
                return None
            
            return model_type(**model_data)
            
        except Exception as e:
            _logger.error(f"Error getting model: {e}")
            return None

    def delete_model(
        self,
        key: str,
        collection: str = "generic_models",
    ) -> bool:
        """Delete a generic model from storage."""
        try:
            coll = self._get_generic_model_collection(collection)
            
            result = coll.delete_one({"_id": key})
            
            if result.deleted_count > 0:
                _logger.debug(f"Deleted model with key '{key}' from collection '{collection}'")
                return True
            return False
            
        except Exception as e:
            _logger.error(f"Error deleting model: {e}")
            return False

    def list_models(
        self,
        model_type: Any,
        collection: str = "generic_models",
    ) -> List[Any]:
        """List all models of a type in a collection."""
        try:
            coll = self._get_generic_model_collection(collection)
            
            cursor = coll.find({"collection": collection})
            
            models = []
            for document in cursor:
                try:
                    model_data = document.get("model_data")
                    if model_data:
                        models.append(model_type(**model_data))
                except Exception as e:
                    _logger.warning(f"Failed to deserialize model: {e}")
            
            return models
            
        except Exception as e:
            _logger.error(f"Error listing models: {e}")
            return []

    # ======================== Cultural Knowledge Methods ========================

    def _get_cultural_knowledge_collection(
        self, create_collection_if_not_found: bool = False
    ) -> Optional["Collection"]:
        """Get the cultural knowledge collection, creating if needed."""
        return self._get_collection(
            "cultural_knowledge",
            create_collection_if_not_found=create_collection_if_not_found,
        )

    def delete_cultural_knowledge(self, id: str) -> None:
        """Delete cultural knowledge from the database.
        
        Args:
            id: The ID of the cultural knowledge to delete.
        
        Raises:
            Exception: If an error occurs during deletion.
        """
        try:
            collection = self._get_cultural_knowledge_collection(
                create_collection_if_not_found=False
            )
            if collection is None:
                return

            collection.delete_one({"id": id})
            _logger.debug(f"Deleted cultural knowledge with ID: {id}")

        except Exception as e:
            _logger.error(f"Error deleting cultural knowledge: {e}")
            raise e

    def get_cultural_knowledge(
        self,
        id: str,
        deserialize: bool = True,
    ) -> Optional[Union["CulturalKnowledge", Dict[str, Any]]]:
        """Get cultural knowledge from the database.
        
        Args:
            id: The ID of the cultural knowledge to get.
            deserialize: If True, return CulturalKnowledge object. If False, return dict.
        
        Returns:
            CulturalKnowledge object or dict if found, None otherwise.
        
        Raises:
            Exception: If an error occurs during retrieval.
        """
        from upsonic.culture.cultural_knowledge import CulturalKnowledge
        
        try:
            collection = self._get_cultural_knowledge_collection(
                create_collection_if_not_found=False
            )
            if collection is None:
                return None

            result = collection.find_one({"id": id})
            if result is None:
                return None

            # Remove MongoDB's _id field
            db_row = remove_mongo_id(result)
            if not deserialize:
                return db_row
            return CulturalKnowledge.from_dict(db_row)

        except Exception as e:
            _logger.error(f"Exception reading from cultural knowledge collection: {e}")
            raise e

    def get_all_cultural_knowledge(
        self,
        name: Optional[str] = None,
        limit: Optional[int] = None,
        page: Optional[int] = None,
        sort_by: Optional[str] = None,
        sort_order: Optional[str] = None,
        agent_id: Optional[str] = None,
        team_id: Optional[str] = None,
        deserialize: bool = True,
    ) -> Union[List["CulturalKnowledge"], Tuple[List[Dict[str, Any]], int]]:
        """Get all cultural knowledge entries from the database.
        
        Args:
            name: Filter by name.
            limit: Maximum number of records to return.
            page: Page number (1-indexed).
            sort_by: Column to sort by.
            sort_order: Sort order ('asc' or 'desc').
            agent_id: Filter by agent ID.
            team_id: Filter by team ID.
            deserialize: If True, return list of CulturalKnowledge objects.
                        If False, return tuple of (list of dicts, total count).
        
        Returns:
            List of CulturalKnowledge objects or tuple of (list of dicts, total count).
        
        Raises:
            Exception: If an error occurs during retrieval.
        """
        from upsonic.culture.cultural_knowledge import CulturalKnowledge
        
        try:
            collection = self._get_cultural_knowledge_collection(
                create_collection_if_not_found=False
            )
            if collection is None:
                return [] if deserialize else ([], 0)

            # Build query
            query: Dict[str, Any] = {}
            if name is not None:
                query["name"] = name
            if agent_id is not None:
                query["agent_id"] = agent_id
            if team_id is not None:
                query["team_id"] = team_id

            # Get total count for pagination
            total_count = collection.count_documents(query)

            # Apply sorting
            sort_criteria = apply_sorting({}, sort_by, sort_order)

            # Build cursor with sorting and pagination
            cursor = collection.find(query)
            if sort_criteria:
                cursor = cursor.sort(sort_criteria)
            
            # Apply pagination
            if limit is not None:
                if page is not None and page > 1:
                    offset = (page - 1) * limit
                    cursor = cursor.skip(offset)
                cursor = cursor.limit(limit)

            # Remove MongoDB's _id field from all results
            db_rows = [remove_mongo_id(item) for item in cursor]

            if not deserialize:
                return db_rows, total_count
            return [CulturalKnowledge.from_dict(row) for row in db_rows]

        except Exception as e:
            _logger.error(f"Error reading from cultural knowledge collection: {e}")
            raise e

    def upsert_cultural_knowledge(
        self,
        cultural_knowledge: "CulturalKnowledge",
        deserialize: bool = True,
    ) -> Optional[Union["CulturalKnowledge", Dict[str, Any]]]:
        """Upsert cultural knowledge into the database.
        
        Args:
            cultural_knowledge: The CulturalKnowledge instance to upsert.
            deserialize: If True, return CulturalKnowledge object. If False, return dict.
        
        Returns:
            The upserted CulturalKnowledge object or dict, or None if error occurs.
        
        Raises:
            Exception: If an error occurs during upsert.
        """
        from upsonic.culture.cultural_knowledge import CulturalKnowledge
        import uuid
        
        try:
            collection = self._get_cultural_knowledge_collection(
                create_collection_if_not_found=True
            )
            if collection is None:
                return None

            if cultural_knowledge.id is None:
                cultural_knowledge.id = str(uuid.uuid4())

            # Use to_dict for serialization
            data = cultural_knowledge.to_dict()
            
            # Convert RFC3339 timestamps to epoch seconds for storage
            if "created_at" in data and data["created_at"] is not None:
                if isinstance(data["created_at"], str):
                    from upsonic.utils.dttm import to_epoch_s
                    data["created_at"] = to_epoch_s(data["created_at"])
            if "updated_at" in data and data["updated_at"] is not None:
                if isinstance(data["updated_at"], str):
                    from upsonic.utils.dttm import to_epoch_s
                    data["updated_at"] = to_epoch_s(data["updated_at"])
            else:
                # Set updated_at to current time if not set
                import time
                data["updated_at"] = int(time.time())

            # Upsert using replace_one
            collection.replace_one(
                {"id": cultural_knowledge.id},
                data,
                upsert=True,
            )

            if not deserialize:
                return data
            return CulturalKnowledge.from_dict(data)

        except Exception as e:
            _logger.error(f"Error upserting cultural knowledge: {e}")
            raise e