"""SessionManager — high-level API for conversation session management."""

from __future__ import annotations

import asyncio
import datetime
import logging
from collections import defaultdict, deque

from agentkit.session.models import Message, MessageRole, Session, SessionStatus
from agentkit.session.store import InMemorySessionStore, SessionStore

logger = logging.getLogger(__name__)


class AsyncWriteQueue:
    """Background pending-buffer queue for non-blocking session persistence.

    Accepts write operations (append_message + save_session) as tasks,
    executes them in a background ``asyncio.Task``, and maintains a small
    in-memory pending buffer for immediate reads.

    WARNING: This is NOT a true Write-Ahead Log. The pending buffer lives
    entirely in memory. If the process crashes before ``flush()`` completes,
    all uncommitted messages will be lost. Only use ``async_writes=True``
    in scenarios where occasional message loss is acceptable.
    """

    def __init__(self, store: SessionStore, max_buffer_size: int = 256) -> None:
        self._store = store
        self._queue: asyncio.Queue[tuple[Message, Session] | None] | None = None
        self._worker: asyncio.Task | None = None
        # Pending buffer: session_id -> deque of Messages not yet persisted
        self._pending_buffer: dict[str, deque[Message]] = defaultdict(deque)
        self._max_buffer_size = max_buffer_size
        self._pending_count = 0

    def _ensure_started(self) -> None:
        """Start the background writer task if not already running (lazy init)."""
        if self._worker is not None and not self._worker.done():
            return
        self._queue = asyncio.Queue()
        self._worker = asyncio.create_task(self._writer_loop())

    async def _writer_loop(self) -> None:
        """Consume write tasks from the queue and persist them."""
        assert self._queue is not None
        while True:
            item = await self._queue.get()
            if item is None:
                # Sentinel: graceful shutdown signal
                self._queue.task_done()
                break
            message, session = item
            try:
                await self._store.append_message(message)
                session.updated_at = datetime.datetime.now(datetime.timezone.utc)
                await self._store.save_session(session)
            except Exception:
                logger.exception(
                    "AsyncWriteQueue: failed to persist message %s for session %s",
                    message.message_id,
                    message.session_id,
                )
            finally:
                # Remove from pending buffer once persisted
                buf = self._pending_buffer.get(message.session_id)
                if buf is not None:
                    try:
                        buf.remove(message)
                    except ValueError:
                        pass
                    if not buf:
                        self._pending_buffer.pop(message.session_id, None)
                self._pending_count = max(0, self._pending_count - 1)
                self._queue.task_done()

    def enqueue(self, message: Message, session: Session) -> None:
        """Enqueue a write task without blocking.

        The message is immediately added to the WAL buffer so that
        ``get_chat_messages`` can see it before persistence completes.
        Lazily starts the background writer on first call.
        """
        self._ensure_started()
        assert self._queue is not None
        # Check buffer size limit before enqueueing
        if len(self._pending_buffer[message.session_id]) >= self._max_buffer_size:
            logger.warning(
                "AsyncWriteQueue: pending buffer full for session %s, flushing before enqueue",
                message.session_id,
            )
            # Synchronous flush to prevent unbounded growth
            self._pending_buffer[message.session_id].append(message)
            self._pending_count += 1
            self._queue.put_nowait((message, session))
            return
        self._pending_buffer[message.session_id].append(message)
        self._pending_count += 1
        self._queue.put_nowait((message, session))

    def buffered_messages(self, session_id: str) -> list[Message]:
        """Return pending-buffered messages for *session_id* not yet persisted."""
        return list(self._pending_buffer.get(session_id, []))

    @property
    def pending_count(self) -> int:
        """Number of write tasks waiting in the queue."""
        return self._pending_count

    async def flush(self) -> None:
        """Wait until all queued writes have been persisted."""
        if self._queue is not None:
            await self._queue.join()

    async def stop(self) -> None:
        """Signal the writer to stop and wait for it to finish."""
        if self._queue is not None and self._worker is not None:
            await self._queue.put(None)  # sentinel
            await self._worker
            self._worker = None
            self._queue = None


class SessionManager:
    """Manages conversation sessions and their messages.

    Provides a high-level API for creating, querying, and updating
    sessions, as well as appending and retrieving messages.

    When ``async_writes=True``, ``append_message`` is non-blocking: the
    message is placed in a write-ahead buffer and persisted in the
    background.  Call ``flush()`` or ``close()`` to ensure all writes
    are durably persisted before shutdown.
    """

    def __init__(
        self,
        store: SessionStore | None = None,
        *,
        async_writes: bool = False,
        wal_buffer_size: int = 256,
    ):
        self._store = store or InMemorySessionStore()
        self._async_writes = async_writes
        self._write_queue: AsyncWriteQueue | None = None
        if async_writes:
            self._write_queue = AsyncWriteQueue(self._store, max_buffer_size=wal_buffer_size)

    @property
    def store(self) -> SessionStore:
        return self._store

    async def create_session(
        self,
        agent_name: str,
        metadata: dict[str, object] | None = None,
    ) -> Session:
        """Create a new conversation session bound to an Agent.

        Args:
            agent_name: Name of the Agent this session is bound to.
            metadata: Optional metadata to attach to the session.

        Returns:
            The newly created Session.
        """
        session = Session(
            session_id=Session.new_session_id(),
            agent_name=agent_name,
            metadata=metadata or {},
        )
        await self._store.save_session(session)
        logger.info(f"Session created: {session.session_id} for agent '{agent_name}'")
        return session

    async def get_session(self, session_id: str) -> Session | None:
        """Get a session by ID."""
        return await self._store.get_session(session_id)

    async def pause_session(self, session_id: str) -> Session | None:
        """Pause an active session."""
        return await self._store.update_session_status(session_id, SessionStatus.PAUSED)

    async def resume_session(self, session_id: str) -> Session | None:
        """Resume a paused session."""
        return await self._store.update_session_status(session_id, SessionStatus.ACTIVE)

    async def close_session(self, session_id: str) -> Session | None:
        """Close a session. Closed sessions cannot accept new messages.

        When async writes are enabled, flushes pending writes before
        updating the session status.
        """
        if self._write_queue is not None:
            await self._write_queue.flush()
        return await self._store.update_session_status(session_id, SessionStatus.CLOSED)

    async def delete_session(self, session_id: str) -> bool:
        """Delete a session and all its messages."""
        return await self._store.delete_session(session_id)

    async def list_sessions(
        self,
        agent_name: str | None = None,
        limit: int = 100,
    ) -> list[Session]:
        """List sessions, optionally filtered by agent name."""
        return await self._store.list_sessions(agent_name=agent_name, limit=limit)

    async def append_message(
        self,
        session_id: str,
        role: MessageRole,
        content: str,
        tool_call_id: str | None = None,
        agent_name: str | None = None,
        metadata: dict[str, object] | None = None,
    ) -> Message:
        """Append a message to a session.

        Args:
            session_id: Target session ID.
            role: Message role (user/assistant/tool/system).
            content: Message content.
            tool_call_id: Optional tool call ID for tool messages.
            agent_name: Optional agent name for multi-Agent sessions.
            metadata: Optional message metadata.

        Returns:
            The newly created Message.

        Raises:
            ValueError: If the session does not exist or is closed.
        """
        session = await self._store.get_session(session_id)
        if session is None:
            raise ValueError(f"Session '{session_id}' not found")
        if session.status == SessionStatus.CLOSED:
            raise ValueError(f"Session '{session_id}' is closed and cannot accept new messages")

        message = Message(
            message_id=Session.new_message_id(),
            session_id=session_id,
            role=role,
            content=content,
            tool_call_id=tool_call_id,
            agent_name=agent_name,
            metadata=metadata or {},
        )

        if self._write_queue is not None:
            # Non-blocking: enqueue write, return immediately
            self._write_queue.enqueue(message, session)
        else:
            # Synchronous path (default, backward-compatible)
            await self._store.append_message(message)
            session.updated_at = datetime.datetime.now(datetime.timezone.utc)
            await self._store.save_session(session)

        return message

    async def get_messages(
        self,
        session_id: str,
        limit: int | None = None,
        offset: int = 0,
    ) -> list[Message]:
        """Get messages for a session with optional pagination.

        When async writes are enabled, includes WAL-buffered messages
        not yet persisted to the store.

        Args:
            session_id: Target session ID.
            limit: Maximum number of messages to return. None for all.
            offset: Number of messages to skip from the beginning.

        Returns:
            List of messages ordered chronologically.
        """
        persisted = await self._store.get_messages(session_id, limit=None, offset=0)
        if self._write_queue is not None:
            buffered = self._write_queue.buffered_messages(session_id)
            if buffered:
                # Merge: persisted + buffered (dedup by message_id)
                seen = {m.message_id for m in persisted}
                for m in buffered:
                    if m.message_id not in seen:
                        persisted.append(m)
        sliced = persisted[offset:]
        if limit is not None:
            sliced = sliced[:limit]
        return sliced

    async def get_chat_messages(self, session_id: str) -> list[dict[str, str]]:
        """Get messages formatted for LLM chat API consumption.

        Returns messages as OpenAI-compatible dicts suitable for
        passing directly to the ReAct engine or LLM Gateway.

        When async writes are enabled, includes WAL-buffered messages
        not yet persisted to the store.
        """
        messages = await self.get_messages(session_id)
        return [m.to_chat_message() for m in messages]

    async def count_messages(self, session_id: str) -> int:
        """Count messages in a session."""
        return await self._store.count_messages(session_id)

    async def health_check(self) -> bool:
        """Check if the underlying store is healthy."""
        return await self._store.health_check()

    async def flush(self) -> None:
        """Wait for all pending async writes to be persisted.

        No-op when async writes are not enabled.
        """
        if self._write_queue is not None:
            await self._write_queue.flush()

    async def close(self) -> None:
        """Flush pending writes and stop the background writer.

        Should be called during application shutdown when async writes
        are enabled.  Safe to call even when async writes are disabled.
        """
        if self._write_queue is not None:
            await self._write_queue.flush()
            await self._write_queue.stop()
