"""
Real-time output streaming system for agent reasoning and output visibility.

This module provides streaming capabilities that allow users to see agent
thinking and output as it happens, not just final results.

Requirements: EXEC-08 (real-time output), EXEC-09 (reasoning visibility)
"""

from typing import Dict, Any, Optional, Callable, Awaitable, AsyncIterator
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
import logging

import anyio
from anyio.streams.memory import MemoryObjectSendStream, MemoryObjectReceiveStream

logger = logging.getLogger(__name__)


class StreamEventType(str, Enum):
    """Types of events that can be streamed during execution.

    Events represent different stages and types of output during
    agent execution, from reasoning to final results.
    """

    REASONING = "reasoning"  # Agent's thinking process
    OUTPUT = "output"  # Final output chunks
    ERROR = "error"  # Error messages
    COMPLETE = "complete"  # Task completed
    PROGRESS = "progress"  # Progress updates (percentage, status)


@dataclass
class StreamEvent:
    """A streaming event emitted during agent execution.

    Events capture the task, event type, content, and timestamp for
    real-time visibility into agent execution.

    Attributes:
        task_id: Identifier of the task generating the event
        event_type: Type of event (reasoning, output, error, etc.)
        content: The content of the event
        timestamp: ISO format timestamp, auto-generated if not provided
        metadata: Optional additional metadata (e.g., progress percentages)
    """

    task_id: str
    event_type: StreamEventType
    content: str
    timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    metadata: Dict[str, Any] = field(default_factory=dict)


class OutputStream:
    """Real-time output streaming for agent execution (EXEC-08, EXEC-09).

    This class provides streaming capabilities using AnyIO memory object
    streams for real-time event distribution. Multiple consumers can
    subscribe to the same task stream.

    Example:
        ```python
        import anyio
        from gsd_rlm.execution.streaming import OutputStream, StreamEventType

        async def main():
            streamer = OutputStream()

            # Create a stream for a task
            receive_stream = await streamer.create_stream("task-1")

            # Emit events
            await streamer.emit("task-1", StreamEventType.REASONING, "Thinking...")
            await streamer.emit("task-1", StreamEventType.OUTPUT, "Result")

            # Receive events
            async for event in receive_stream:
                print(f"[{event.event_type.value}] {event.content}")
                if event.event_type == StreamEventType.COMPLETE:
                    break

            # Clean up
            await streamer.close_stream("task-1")

        anyio.run(main)
        ```
    """

    def __init__(self, buffer_size: int = 100):
        """Initialize the output stream.

        Args:
            buffer_size: Maximum number of events to buffer per stream.
                        Default 100 provides good balance of memory vs throughput.
        """
        self.buffer_size = buffer_size
        self._streams: Dict[str, MemoryObjectSendStream] = {}

    async def create_stream(self, task_id: str) -> MemoryObjectReceiveStream:
        """Create a new stream for a task.

        Creates a send/receive stream pair and stores the send stream
        internally for emitting events. The receive stream is returned
        for the consumer to iterate over.

        Args:
            task_id: Unique identifier for the task

        Returns:
            MemoryObjectReceiveStream for receiving StreamEvent objects

        Raises:
            ValueError: If a stream already exists for this task_id
        """
        if task_id in self._streams:
            raise ValueError(f"Stream already exists for task_id: {task_id}")

        send_stream, receive_stream = anyio.create_memory_object_stream[StreamEvent](
            max_buffer_size=self.buffer_size
        )
        self._streams[task_id] = send_stream
        return receive_stream

    async def emit(
        self,
        task_id: str,
        event_type: StreamEventType,
        content: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Emit a stream event for real-time visibility.

        Creates a StreamEvent with auto-generated timestamp and sends
        it to the stored send stream for the task.

        Args:
            task_id: Identifier of the task
            event_type: Type of event to emit
            content: Content of the event
            metadata: Optional additional metadata

        Note:
            Silently does nothing if no stream exists for the task_id.
            This allows safe emission even if consumer disconnected.
        """
        if task_id not in self._streams:
            logger.debug(f"No stream for task_id {task_id}, skipping emit")
            return

        event = StreamEvent(
            task_id=task_id,
            event_type=event_type,
            content=content,
            timestamp=datetime.utcnow().isoformat(),
            metadata=metadata or {},
        )

        await self._streams[task_id].send(event)
        logger.debug(f"Emitted {event_type.value} event for task {task_id}")

    async def close_stream(self, task_id: str) -> None:
        """Close and remove the stream for a task.

        Closes the send stream and removes it from internal storage.
        Should be called when task execution completes or fails.

        Args:
            task_id: Identifier of the task

        Note:
            Silently does nothing if no stream exists for the task_id.
        """
        if task_id in self._streams:
            await self._streams[task_id].aclose()
            del self._streams[task_id]
            logger.debug(f"Closed stream for task {task_id}")

    async def stream_task(
        self,
        task_id: str,
        executor: Callable[[str], Awaitable[Any]],
    ) -> AsyncIterator[StreamEvent]:
        """Execute a task and yield stream events in real-time.

        Convenience method that handles the full stream lifecycle:
        creates stream, runs executor in background, yields events
        as they arrive, and emits COMPLETE event when done.

        Args:
            task_id: Unique identifier for the task
            executor: Async callable that executes the task. The executor
                     should use the OutputStream to emit events during execution.

        Yields:
            StreamEvent objects as they are emitted

        Example:
            ```python
            async def my_executor(task_id: str):
                await streamer.emit(task_id, StreamEventType.REASONING, "Analyzing...")
                # Do work...
                await streamer.emit(task_id, StreamEventType.OUTPUT, "Result")
                return "done"

            async for event in streamer.stream_task("task-1", my_executor):
                print(f"[{event.event_type.value}] {event.content}")
            ```
        """
        receive_stream = await self.create_stream(task_id)
        execution_complete = False

        async def execute_and_emit():
            """Run executor and emit completion/error events."""
            nonlocal execution_complete
            try:
                await executor(task_id)
                # Emit completion event if executor didn't
                if not execution_complete:
                    await self.emit(task_id, StreamEventType.COMPLETE, "")
            except Exception as e:
                await self.emit(task_id, StreamEventType.ERROR, str(e))
                await self.emit(task_id, StreamEventType.COMPLETE, "")
            finally:
                await self.close_stream(task_id)

        # Run execution in background task group
        async with anyio.create_task_group() as tg:
            tg.start_soon(execute_and_emit)

            # Yield events as they arrive
            async for event in receive_stream:
                yield event
                if event.event_type == StreamEventType.COMPLETE:
                    execution_complete = True
                    break

    def has_stream(self, task_id: str) -> bool:
        """Check if a stream exists for the given task.

        Args:
            task_id: Identifier of the task

        Returns:
            True if stream exists, False otherwise
        """
        return task_id in self._streams

    def active_streams(self) -> list:
        """Get list of active stream task IDs.

        Returns:
            List of task_ids that have active streams
        """
        return list(self._streams.keys())
