"""Session management helpers for Agentbyte WebUI."""

from __future__ import annotations

import uuid

from agentbyte.context import AgentContext

from .models import SessionInfo
from .session_store import InMemorySessionStore, SessionStore


class SessionManager:
    """Manage AgentContext-backed UI sessions."""

    def __init__(self, store: SessionStore | None = None) -> None:
        self.store = store or InMemorySessionStore()

    def create_session_id(self) -> str:
        """Return a fresh session id."""
        return str(uuid.uuid4())

    async def get_or_create(
        self,
        session_id: str,
        entity_id: str,
        entity_type: str,
    ) -> AgentContext:
        """Return an existing session or create a new empty context."""
        context = await self.store.get(session_id)
        if context is None:
            context = AgentContext(
                session_id=session_id,
                metadata={
                    "entity_id": entity_id,
                    "entity_type": entity_type,
                },
            )
            await self.store.save(session_id, context)
        return context

    async def get(self, session_id: str) -> AgentContext | None:
        """Load a session context by id."""
        return await self.store.get(session_id)

    async def update(self, session_id: str, context: AgentContext) -> None:
        """Persist an updated session context."""
        await self.store.save(session_id, context)

    async def delete(self, session_id: str) -> bool:
        """Delete a stored session."""
        return await self.store.delete(session_id)

    async def clear_all(self) -> int:
        """Clear all sessions."""
        return await self.store.clear_all()

    async def list(self, entity_id: str | None = None) -> list[SessionInfo]:
        """List session metadata, optionally filtered by entity id."""
        sessions: list[SessionInfo] = []
        for session_id, context in await self.store.list():
            metadata_entity_id = str(context.metadata.get("entity_id", ""))
            if entity_id is not None and metadata_entity_id != entity_id:
                continue

            sessions.append(
                SessionInfo(
                    id=session_id,
                    entity_id=metadata_entity_id,
                    entity_type=str(context.metadata.get("entity_type", "agent")),
                    created_at=context.created_at,
                    last_activity=context.updated_at or context.created_at,
                    message_count=len(context.messages),
                    user_id=context.user_id,
                )
            )

        sessions.sort(key=lambda item: item.last_activity, reverse=True)
        return sessions


__all__ = ["SessionManager"]
