steering.py

"""
Steering - Streaming conversation mode for queuing messages with priority support.

This module provides the infrastructure for queuing multiple user messages
while the CLI processes previous ones, with support for priority-based ordering
and graceful shutdown.

Example:
    >>> from copilot import CopilotClient
    >>> from copilot.steering import ConversationManager, Priority
    >>>
    >>> async with CopilotClient() as client:
    ...     session = await client.create_session()
    ...     manager = ConversationManager(session)
    ...
    ...     # Queue messages - returns immediately even if CLI is busy
    ...     await manager.queue_message("req-1", "What is Python?", Priority.NORMAL)
    ...     await manager.queue_message("req-2", "URGENT: Fix the bug!", Priority.URGENT)
    ...     await manager.queue_message("req-3", "Tell me a joke", Priority.LOW)
    ...
    ...     # Messages processed in order: req-2, req-1, req-3
    ...
    ...     # Graceful shutdown
    ...     await manager.stop()
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from datetime import datetime
from enum import IntEnum
from typing import Any, AsyncIterator, Callable, Optional

from .session import CopilotSession
from .types import Attachment


class Priority(IntEnum):
    """Message priority levels for queue ordering.

    Higher values are processed first. Within the same priority,
    messages are processed in FIFO order.
    """

    LOW = 0
    NORMAL = 1
    HIGH = 2
    URGENT = 3


class QueueFullError(Exception):
    """Raised when the message queue is full and cannot accept new messages."""

    pass


class ShutdownSentinel:
    """Marker to signal generator termination.

    Always sorts last in the priority queue to ensure all pending
    messages are processed before shutdown.
    """

    def __lt__(self, other: object) -> bool:
        """Always sort after real messages."""
        return False

    def __le__(self, other: object) -> bool:
        """Always sort after real messages."""
        return isinstance(other, ShutdownSentinel)

    def __gt__(self, other: object) -> bool:
        """Always sort after real messages."""
        return not isinstance(other, ShutdownSentinel)

    def __ge__(self, other: object) -> bool:
        """Always sort after real messages."""
        return True


# Singleton sentinel instance
SHUTDOWN_SENTINEL = ShutdownSentinel()


@dataclass(order=False)
class QueuedMessage:
    """Represents a message queued for processing.

    Attributes:
        request_id: Unique identifier for this message request.
        content: The message content/prompt text.
        priority: Processing priority (higher = processed first).
        session_id: The session this message belongs to.
        queued_at: When the message was queued.
        metadata: Additional metadata for the message.
        sequence_number: For FIFO ordering within same priority.
        attachments: Optional file/directory attachments.
    """

    request_id: str
    content: str
    priority: Priority
    session_id: str
    queued_at: datetime = field(default_factory=datetime.now)
    metadata: dict[str, Any] = field(default_factory=dict)
    sequence_number: int = 0
    attachments: list[Attachment] = field(default_factory=list)

    def __lt__(self, other: object) -> bool:
        """Priority queue ordering: higher priority first, then FIFO."""
        if isinstance(other, ShutdownSentinel):
            return True  # Real messages come before sentinel
        if not isinstance(other, QueuedMessage):
            return NotImplemented
        if self.priority != other.priority:
            return self.priority.value > other.priority.value  # Higher = first
        return self.sequence_number < other.sequence_number  # Earlier = first

    def __le__(self, other: object) -> bool:
        if isinstance(other, ShutdownSentinel):
            return True
        if not isinstance(other, QueuedMessage):
            return NotImplemented
        return self < other or (
            self.priority == other.priority
            and self.sequence_number == other.sequence_number
        )

    def __gt__(self, other: object) -> bool:
        if isinstance(other, ShutdownSentinel):
            return False
        if not isinstance(other, QueuedMessage):
            return NotImplemented
        return not self <= other

    def __ge__(self, other: object) -> bool:
        if isinstance(other, ShutdownSentinel):
            return False
        if not isinstance(other, QueuedMessage):
            return NotImplemented
        return not self < other


class MessageQueue:
    """Priority queue for conversation messages.

    Provides non-blocking put() and blocking get() semantics with
    priority-based ordering. Thread-safe for use with asyncio.

    Attributes:
        max_depth: Maximum number of messages the queue can hold.
    """

    def __init__(self, max_depth: int = 100):
        """Initialize the message queue.

        Args:
            max_depth: Maximum queue size. Defaults to 100.
        """
        self._max_depth = max_depth
        self._queue: asyncio.PriorityQueue[QueuedMessage | ShutdownSentinel] = (
            asyncio.PriorityQueue(maxsize=max_depth)
        )
        self._shutdown_event = asyncio.Event()
        self._sequence_counter = 0

    @property
    def max_depth(self) -> int:
        """Maximum queue capacity."""
        return self._max_depth

    def qsize(self) -> int:
        """Return the current queue size."""
        return self._queue.qsize()

    def empty(self) -> bool:
        """Return True if the queue is empty."""
        return self._queue.empty()

    def full(self) -> bool:
        """Return True if the queue is full."""
        return self._queue.full()

    def _next_sequence(self) -> int:
        """Get the next sequence number for FIFO ordering."""
        seq = self._sequence_counter
        self._sequence_counter += 1
        return seq

    async def put(self, message: QueuedMessage) -> None:
        """Add a message to the queue.

        This is non-blocking - it raises QueueFullError if the queue is full
        rather than waiting.

        Args:
            message: The message to queue.

        Raises:
            QueueFullError: If the queue is at max capacity.
        """
        # Assign sequence number if not set
        if message.sequence_number == 0:
            message.sequence_number = self._next_sequence()

        try:
            self._queue.put_nowait(message)
        except asyncio.QueueFull:
            raise QueueFullError(f"Queue full (max={self._max_depth})")

    async def get(self) -> QueuedMessage | ShutdownSentinel:
        """Get the next message from the queue.

        Blocks until a message is available or shutdown is signaled.

        Returns:
            The next message to process, or ShutdownSentinel if shutting down.
        """
        return await self._queue.get()

    def signal_shutdown(self) -> None:
        """Signal the generator to terminate.

        This unblocks any waiting get() calls and causes the generator
        to terminate gracefully.
        """
        self._shutdown_event.set()
        try:
            self._queue.put_nowait(SHUTDOWN_SENTINEL)
        except asyncio.QueueFull:
            # Queue is full, but we still set the event
            pass

    def is_shutdown(self) -> bool:
        """Check if shutdown has been signaled."""
        return self._shutdown_event.is_set()


class StreamingInputGenerator:
    """Async generator that yields messages from queue to SDK.

    This class bridges the message queue to the SDK's async iterator
    interface, yielding formatted message dicts until shutdown.

    Example:
        >>> queue = MessageQueue()
        >>> generator = StreamingInputGenerator(queue)
        >>> async for message in generator:
        ...     # Process message
        ...     pass
    """

    def __init__(self, queue: MessageQueue):
        """Initialize the generator.

        Args:
            queue: The message queue to consume from.
        """
        self._queue = queue

    def __aiter__(self) -> AsyncIterator[dict[str, Any]]:
        """Return self as the async iterator."""
        return self._generate()

    async def _generate(self) -> AsyncIterator[dict[str, Any]]:
        """Yield messages until shutdown sentinel received."""
        while True:
            item = await self._queue.get()

            if isinstance(item, ShutdownSentinel):
                break

            # Format for SDK consumption
            yield {
                "type": "user",
                "message": {
                    "role": "user",
                    "content": item.content,
                },
                "metadata": {
                    "request_id": item.request_id,
                    "priority": item.priority.name,
                    "session_id": item.session_id,
                    "queued_at": item.queued_at.isoformat(),
                    **item.metadata,
                },
                "attachments": item.attachments,
            }


class ConversationManager:
    """Orchestrates message flow from callers to SDK session.

    The ConversationManager provides a high-level interface for queuing
    messages while the CLI processes previous ones. It handles:

    - Non-blocking message queuing with priority support
    - Automatic session interaction
    - Graceful shutdown

    Example:
        >>> async with CopilotClient() as client:
        ...     session = await client.create_session()
        ...     manager = ConversationManager(session)
        ...
        ...     # Queue messages - returns immediately
        ...     await manager.queue_message("req-1", "Hello", Priority.NORMAL)
        ...     await manager.queue_message("req-2", "Urgent!", Priority.URGENT)
        ...
        ...     # Stop when done
        ...     await manager.stop()
    """

    def __init__(
        self,
        session: CopilotSession,
        max_queue_depth: int = 100,
        on_response: Optional[Callable[[dict[str, Any]], None]] = None,
    ):
        """Initialize the conversation manager.

        Args:
            session: The CopilotSession to send messages to.
            max_queue_depth: Maximum number of queued messages. Defaults to 100.
            on_response: Optional callback for responses/events from the session.
        """
        self._session = session
        self._max_queue_depth = max_queue_depth
        self._on_response = on_response
        self._queue: Optional[MessageQueue] = None
        self._generator: Optional[StreamingInputGenerator] = None
        self._processor_task: Optional[asyncio.Task[None]] = None
        self._started = False
        self._request_counter = 0

    @property
    def session(self) -> CopilotSession:
        """The underlying session."""
        return self._session

    @property
    def is_started(self) -> bool:
        """Whether the manager has been started."""
        return self._started

    @property
    def queue_size(self) -> int:
        """Current number of queued messages."""
        return self._queue.qsize() if self._queue else 0

    def _generate_request_id(self) -> str:
        """Generate a unique request ID."""
        self._request_counter += 1
        return f"req-{self._request_counter}"

    async def queue_message(
        self,
        content: str,
        priority: Priority = Priority.NORMAL,
        request_id: Optional[str] = None,
        attachments: Optional[list[Attachment]] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> str:
        """Queue a message for processing.

        This method returns immediately, even if the CLI is busy processing
        a previous message. Messages are processed in priority order.

        Args:
            content: The message content/prompt.
            priority: Processing priority. Defaults to NORMAL.
            request_id: Optional custom request ID. Auto-generated if not provided.
            attachments: Optional file/directory attachments.
            metadata: Optional additional metadata.

        Returns:
            The request ID for this message.

        Raises:
            QueueFullError: If the queue is at max capacity.
        """
        # Auto-start on first message
        if not self._started:
            await self._start()

        if request_id is None:
            request_id = self._generate_request_id()

        msg = QueuedMessage(
            request_id=request_id,
            content=content,
            priority=priority,
            session_id=self._session.session_id,
            attachments=attachments or [],
            metadata=metadata or {},
        )
        await self._queue.put(msg)  # type: ignore[union-attr]
        return request_id

    async def _start(self) -> None:
        """Start the conversation processing loop."""
        if self._started:
            return

        self._queue = MessageQueue(max_depth=self._max_queue_depth)
        self._generator = StreamingInputGenerator(self._queue)
        self._started = True

        # Start the processor task that consumes from the generator
        self._processor_task = asyncio.create_task(self._process_messages())

    async def _process_messages(self) -> None:
        """Process messages from the generator, sending them to the session."""
        if self._generator is None:
            return

        async for message_dict in self._generator:
            try:
                # Extract message details
                content = message_dict["message"]["content"]
                attachments = message_dict.get("attachments", [])
                request_id = message_dict["metadata"]["request_id"]

                # Send to session - this is non-blocking in the SDK
                # The session.send returns a message ID
                await self._session.send(
                    {
                        "prompt": content,
                        "attachments": attachments if attachments else None,
                    }
                )

                # If we have a response callback, we could wire it up here
                # For now, the caller should use session.on() for events

            except Exception as e:
                # Log error but continue processing
                print(f"Error processing message {message_dict.get('metadata', {}).get('request_id', 'unknown')}: {e}")

    async def stop(self, timeout: Optional[float] = None) -> None:
        """Gracefully stop the conversation manager.

        Signals shutdown and waits for pending messages to be processed.

        Args:
            timeout: Maximum time to wait for pending messages. Defaults to None (wait forever).
        """
        if not self._started or self._queue is None:
            return

        # Signal shutdown
        self._queue.signal_shutdown()

        # Wait for processor task to complete
        if self._processor_task is not None:
            try:
                if timeout is not None:
                    await asyncio.wait_for(self._processor_task, timeout=timeout)
                else:
                    await self._processor_task
            except asyncio.TimeoutError:
                self._processor_task.cancel()
                try:
                    await self._processor_task
                except asyncio.CancelledError:
                    pass
            except asyncio.CancelledError:
                pass

        self._started = False
        self._queue = None
        self._generator = None
        self._processor_task = None

    async def clear_queue(self) -> int:
        """Clear all pending messages from the queue.

        Returns:
            The number of messages that were cleared.
        """
        if self._queue is None:
            return 0

        count = 0
        while not self._queue.empty():
            try:
                # We need to drain the queue
                # Create a new empty queue and swap
                item = self._queue._queue.get_nowait()
                if not isinstance(item, ShutdownSentinel):
                    count += 1
            except asyncio.QueueEmpty:
                break

        return count

    async def __aenter__(self) -> "ConversationManager":
        """Async context manager entry."""
        return self

    async def __aexit__(
        self,
        exc_type: Optional[type],
        exc_val: Optional[BaseException],
        exc_tb: Optional[Any],
    ) -> None:
        """Async context manager exit - ensures clean shutdown."""
        await self.stop()