import asyncio
from typing import Dict, Optional
from agent.memory import MemoryManager


class SessionManager:
    """Manages agent sessions and their associated memory.

    Coordinates between the gateway's UserSession (for conversation state)
    and the agent's MemoryManager (for persistent memory with Zep).
    """

    def __init__(self, memory_manager: MemoryManager):
        """Initialize session manager.

        Args:
            memory_manager: MemoryManager instance for persistent storage
        """
        self.memory = memory_manager
        self._locks: Dict[str, asyncio.Lock] = {}

    async def get_session_lock(self, session_id: str) -> asyncio.Lock:
        """Get or create a lock for a session.

        Prevents concurrent execution for the same session.

        Args:
            session_id: Session identifier

        Returns:
            asyncio.Lock for this session
        """
        if session_id not in self._locks:
            self._locks[session_id] = asyncio.Lock()
        return self._locks[session_id]

    async def save_message(
        self,
        session_id: str,
        role: str,
        content: str,
        metadata: Optional[Dict] = None
    ) -> None:
        """Save a message to persistent memory.

        Args:
            session_id: Session identifier
            role: Message role ("user" or "assistant")
            content: Message content
            metadata: Optional metadata
        """
        await self.memory.add_memory(session_id, role, content, metadata)

    async def get_relevant_context(
        self,
        session_id: str,
        query: str,
        limit: int = 5
    ) -> str:
        """Get relevant historical context for a query.

        Uses semantic search over past conversations.

        Args:
            session_id: Session identifier
            query: Search query
            limit: Maximum results

        Returns:
            Formatted string of relevant past messages
        """
        results = await self.memory.search_memory(session_id, query, limit)

        if not results:
            return ""

        context_parts = ["## Relevant Past Context\n"]
        for i, result in enumerate(results, 1):
            role = result["role"].capitalize()
            content = result["content"]
            score = result.get("score", 0.0)
            context_parts.append(f"{i}. [{role}, relevance: {score:.2f}] {content}\n")

        return "\n".join(context_parts)

    async def cleanup_session(self, session_id: str) -> None:
        """Cleanup session resources.

        Args:
            session_id: Session to cleanup
        """
        if session_id in self._locks:
            del self._locks[session_id]