"""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 SessionBindingError(RuntimeError):
    """Raised when an existing session is reused with incompatible bindings."""


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,
        user_id: str | None = None,
    ) -> 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,
                user_id=user_id,
                metadata={
                    "entity_id": entity_id,
                    "entity_type": entity_type,
                },
            )
            await self.store.save(session_id, context)
            return context

        self._ensure_binding(context, entity_id=entity_id, user_id=user_id)

        needs_save = False
        if context.session_id != session_id:
            context.session_id = session_id
            needs_save = True
        if context.user_id is None and user_id is not None:
            context.user_id = user_id
            needs_save = True
        if "entity_id" not in context.metadata:
            context.metadata["entity_id"] = entity_id
            needs_save = True
        if "entity_type" not in context.metadata:
            context.metadata["entity_type"] = entity_type
            needs_save = True
        if needs_save:
            await self.store.save(session_id, context)
        return context

    async def assert_compatible(
        self,
        session_id: str | None,
        entity_id: str,
        user_id: str | None = None,
    ) -> None:
        """Raise if an existing session is not compatible with the requested bindings."""
        if session_id is None:
            return

        context = await self.store.get(session_id)
        if context is None:
            return

        self._ensure_binding(context, entity_id=entity_id, user_id=user_id)

    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,
        user_id: str | None = None,
    ) -> list[SessionInfo]:
        """List session metadata, optionally filtered by entity and user id."""
        sessions: list[SessionInfo] = []
        session_items = (
            await self.store.list_by_user(user_id) if user_id is not None else await self.store.list()
        )
        for session_id, context in session_items:
            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

    def _ensure_binding(
        self,
        context: AgentContext,
        entity_id: str,
        user_id: str | None,
    ) -> None:
        stored_entity_id = str(context.metadata.get("entity_id", "") or "")
        if stored_entity_id and stored_entity_id != entity_id:
            raise SessionBindingError(
                f"Session {context.session_id} is bound to entity {stored_entity_id}, not {entity_id}"
            )

        if context.user_id is not None and user_id is not None and context.user_id != user_id:
            raise SessionBindingError(
                f"Session {context.session_id} is owned by user {context.user_id}, not {user_id}"
            )


__all__ = ["SessionBindingError", "SessionManager"]
