"""Stdio interface for communication with Kiro.

This module implements the stdio MCP server interface that Kiro expects,
handling async stdin/stdout communication, subprocess lifecycle management,
message queuing, and signal handling.

Requirements satisfied:
- 1.1: Present valid stdio MCP server interface to Kiro
- 1.4: Maintain message ordering and correlation between stdio and HTTP protocols
"""

import asyncio
import signal
import sys
from typing import Optional, Callable, Awaitable, Union
from asyncio import Queue, StreamReader, StreamWriter

from ..models import JsonRpcMessage
from ..exceptions import ProtocolError, InvalidJsonRpcError
from .jsonrpc import JsonRpcParser, JsonRpcSerializer
from ..logging_config import get_logger

logger = get_logger(__name__)

# Default queue size limit to prevent unbounded memory growth
DEFAULT_QUEUE_SIZE = 1000

# Timeout for graceful shutdown operations
SHUTDOWN_TIMEOUT = 5.0

# Timeout for draining pending messages on shutdown
DRAIN_TIMEOUT = 2.0


class StdioInterface:
    """Stdio interface for JSON-RPC communication with Kiro.
    
    This class manages async stdin/stdout communication, providing a clean
    interface for sending and receiving JSON-RPC messages. It handles:
    - Async reading from stdin with proper buffering
    - Async writing to stdout with message queuing
    - Subprocess lifecycle and signal management
    - Graceful shutdown and cleanup
    
    The interface maintains message ordering by using FIFO queues for both
    incoming and outgoing messages, ensuring that messages are processed
    and sent in the order they were received/queued.
    
    Attributes:
        is_running: Whether the interface is currently running
        pending_incoming: Number of messages waiting to be processed
        pending_outgoing: Number of messages waiting to be sent
    """

    def __init__(
        self,
        message_handler: Optional[Callable[[JsonRpcMessage], Awaitable[Optional[JsonRpcMessage]]]] = None,
        queue_size: int = DEFAULT_QUEUE_SIZE
    ):
        """Initialize the stdio interface.
        
        Args:
            message_handler: Optional async callback for handling incoming messages.
                           Should accept a JsonRpcMessage and return a JsonRpcMessage response
                           or None for notifications.
            queue_size: Maximum size for message queues (default: 1000)
        """
        self._message_handler = message_handler
        self._running = False
        self._shutdown_event = asyncio.Event()
        self._started_event = asyncio.Event()
        
        # Message queues for async processing with size limits
        self._incoming_queue: Queue[JsonRpcMessage] = Queue(maxsize=queue_size)
        self._outgoing_queue: Queue[JsonRpcMessage] = Queue(maxsize=queue_size)
        
        # Stream readers/writers
        self._stdin_reader: Optional[StreamReader] = None
        self._stdout_writer: Optional[StreamWriter] = None
        
        # Background tasks
        self._tasks: list[asyncio.Task] = []
        
        # Parser and serializer
        self._parser = JsonRpcParser()
        self._serializer = JsonRpcSerializer()
        
        # Track if we're in the middle of shutdown
        self._shutting_down = False
        
        logger.info("StdioInterface initialized", queue_size=queue_size)

    @property
    def is_running(self) -> bool:
        """Check if the interface is currently running."""
        return self._running

    @property
    def pending_incoming(self) -> int:
        """Get the number of messages waiting to be processed."""
        return self._incoming_queue.qsize()

    @property
    def pending_outgoing(self) -> int:
        """Get the number of messages waiting to be sent."""
        return self._outgoing_queue.qsize()

    def set_message_handler(
        self,
        handler: Callable[[JsonRpcMessage], Awaitable[Optional[JsonRpcMessage]]]
    ) -> None:
        """Set or update the message handler.
        
        Args:
            handler: Async callback for handling incoming messages.
                    Should return a JsonRpcMessage for requests or None for notifications.
        """
        self._message_handler = handler
        logger.debug("Message handler updated")

    async def start(self) -> None:
        """Start the stdio interface and begin processing messages.
        
        This method:
        1. Sets up stdin/stdout streams
        2. Installs signal handlers for graceful shutdown
        3. Starts background tasks for reading, writing, and processing
        
        Raises:
            RuntimeError: If interface is already running
        """
        if self._running:
            raise RuntimeError("StdioInterface is already running")
        
        logger.info("Starting StdioInterface")
        self._running = True
        self._shutting_down = False
        self._shutdown_event.clear()
        self._started_event.clear()
        
        # Set up stdin/stdout streams
        await self._setup_streams()
        
        # Install signal handlers for graceful shutdown
        self._install_signal_handlers()
        
        # Start background tasks
        self._tasks = [
            asyncio.create_task(self._read_loop(), name="stdio-read"),
            asyncio.create_task(self._write_loop(), name="stdio-write"),
            asyncio.create_task(self._process_loop(), name="stdio-process"),
        ]
        
        # Signal that we're fully started
        self._started_event.set()
        
        logger.info("StdioInterface started successfully")

    async def wait_until_started(self, timeout: Optional[float] = None) -> bool:
        """Wait until the interface has fully started.
        
        Args:
            timeout: Maximum time to wait in seconds (None for no timeout)
            
        Returns:
            True if started successfully, False if timeout occurred
        """
        try:
            if timeout is not None:
                await asyncio.wait_for(self._started_event.wait(), timeout=timeout)
            else:
                await self._started_event.wait()
            return True
        except asyncio.TimeoutError:
            return False

    async def stop(self) -> None:
        """Stop the stdio interface and cleanup resources.
        
        This method:
        1. Signals shutdown to all background tasks
        2. Drains pending outgoing messages
        3. Waits for tasks to complete gracefully
        4. Closes streams and cleans up resources
        """
        if not self._running:
            logger.warning("StdioInterface is not running")
            return
        
        if self._shutting_down:
            logger.warning("StdioInterface is already shutting down")
            return
        
        logger.info("Stopping StdioInterface")
        self._shutting_down = True
        self._running = False
        self._shutdown_event.set()
        
        # Drain pending outgoing messages before stopping
        await self._drain_outgoing_queue()
        
        # Cancel all background tasks
        for task in self._tasks:
            if not task.done():
                task.cancel()
        
        # Wait for tasks to complete (with timeout)
        if self._tasks:
            try:
                await asyncio.wait_for(
                    asyncio.gather(*self._tasks, return_exceptions=True),
                    timeout=SHUTDOWN_TIMEOUT
                )
            except asyncio.TimeoutError:
                logger.warning("Timeout waiting for tasks to complete")
        
        # Close streams
        await self._close_streams()
        
        # Clear queues
        self._clear_queues()
        
        self._tasks.clear()
        self._shutting_down = False
        logger.info("StdioInterface stopped")

    async def _drain_outgoing_queue(self) -> None:
        """Drain and send all pending outgoing messages before shutdown.
        
        This ensures that any queued responses are sent before the interface
        stops, preventing message loss.
        """
        if not self._stdout_writer or self._outgoing_queue.empty():
            return
        
        logger.info("Draining outgoing message queue", pending=self._outgoing_queue.qsize())
        
        try:
            async with asyncio.timeout(DRAIN_TIMEOUT):
                while not self._outgoing_queue.empty():
                    try:
                        message = self._outgoing_queue.get_nowait()
                        line = self._serializer.serialize(message)
                        self._stdout_writer.write(line.encode('utf-8'))
                        await self._stdout_writer.drain()
                        logger.debug("Drained message", message_id=message.id)
                    except asyncio.QueueEmpty:
                        break
        except asyncio.TimeoutError:
            logger.warning("Timeout draining outgoing queue", remaining=self._outgoing_queue.qsize())

    async def send_message(self, message: JsonRpcMessage) -> None:
        """Send a JSON-RPC message to stdout.
        
        Messages are queued and sent asynchronously in FIFO order,
        maintaining message ordering as required by the MCP protocol.
        
        Args:
            message: JsonRpcMessage to send
            
        Raises:
            RuntimeError: If interface is not running
            asyncio.QueueFull: If the outgoing queue is full
        """
        if not self._running and not self._shutting_down:
            raise RuntimeError("StdioInterface is not running")
        
        try:
            # Use put_nowait to avoid blocking if queue is full
            self._outgoing_queue.put_nowait(message)
            logger.debug("Queued outgoing message", message_id=message.id)
        except asyncio.QueueFull:
            logger.error("Outgoing queue is full, message dropped", message_id=message.id)
            raise

    async def receive_message(self, timeout: Optional[float] = None) -> Optional[JsonRpcMessage]:
        """Receive a JSON-RPC message from stdin.
        
        Args:
            timeout: Maximum time to wait in seconds (None for no timeout)
        
        Returns:
            JsonRpcMessage received from stdin, or None if timeout occurred
            
        Raises:
            RuntimeError: If interface is not running
        """
        if not self._running:
            raise RuntimeError("StdioInterface is not running")
        
        try:
            if timeout is not None:
                message = await asyncio.wait_for(
                    self._incoming_queue.get(),
                    timeout=timeout
                )
            else:
                message = await self._incoming_queue.get()
            logger.debug("Received message from queue", message_id=message.id)
            return message
        except asyncio.TimeoutError:
            return None

    async def wait_for_shutdown(self) -> None:
        """Wait for shutdown signal.
        
        This method blocks until the interface is signaled to shut down,
        useful for keeping the main process alive.
        """
        await self._shutdown_event.wait()

    # ========================================================================
    # Private Methods - Stream Setup
    # ========================================================================

    async def _setup_streams(self) -> None:
        """Set up async stdin/stdout streams."""
        loop = asyncio.get_event_loop()
        
        # Create stdin reader
        self._stdin_reader = asyncio.StreamReader()
        protocol = asyncio.StreamReaderProtocol(self._stdin_reader)
        await loop.connect_read_pipe(lambda: protocol, sys.stdin)
        
        # Create stdout writer
        transport, protocol = await loop.connect_write_pipe(
            asyncio.streams.FlowControlMixin,
            sys.stdout
        )
        self._stdout_writer = asyncio.StreamWriter(
            transport, protocol, None, loop
        )
        
        logger.debug("Stdio streams set up successfully")

    async def _close_streams(self) -> None:
        """Close stdin/stdout streams."""
        if self._stdout_writer:
            try:
                self._stdout_writer.close()
                await self._stdout_writer.wait_closed()
            except Exception as e:
                logger.error(f"Error closing stdout writer: {e}")
        
        self._stdin_reader = None
        self._stdout_writer = None
        logger.debug("Stdio streams closed")

    # ========================================================================
    # Private Methods - Signal Handling
    # ========================================================================

    def _install_signal_handlers(self) -> None:
        """Install signal handlers for graceful shutdown."""
        loop = asyncio.get_event_loop()
        
        for sig in (signal.SIGTERM, signal.SIGINT):
            try:
                loop.add_signal_handler(
                    sig,
                    lambda s=sig: asyncio.create_task(self._handle_signal(s))
                )
                logger.debug(f"Installed signal handler for {sig.name}")
            except NotImplementedError:
                # Signal handlers not supported on this platform (e.g., Windows)
                logger.warning(f"Signal handlers not supported for {sig.name}")

    async def _handle_signal(self, sig: signal.Signals) -> None:
        """Handle shutdown signals.
        
        Args:
            sig: Signal that was received
        """
        logger.info(f"Received signal {sig.name}, initiating shutdown")
        await self.stop()

    # ========================================================================
    # Private Methods - Background Tasks
    # ========================================================================

    async def _read_loop(self) -> None:
        """Background task for reading messages from stdin.
        
        Continuously reads newline-delimited JSON-RPC messages from stdin
        and queues them for processing. Maintains message ordering by
        processing messages in the order they are received.
        """
        logger.info("Read loop started")
        
        try:
            while self._running and self._stdin_reader:
                try:
                    # Read a line from stdin (newline-delimited)
                    line_bytes = await self._stdin_reader.readline()
                    
                    if not line_bytes:
                        # EOF reached - initiate graceful shutdown
                        logger.info("EOF reached on stdin, initiating shutdown")
                        asyncio.create_task(self.stop())
                        break
                    
                    # Decode and parse the message
                    line = line_bytes.decode('utf-8')
                    
                    # Skip empty lines
                    if not line.strip():
                        continue
                    
                    try:
                        message = self._parser.parse(line)
                        try:
                            self._incoming_queue.put_nowait(message)
                            logger.debug("Read message from stdin", message_id=message.id, method=message.method)
                        except asyncio.QueueFull:
                            logger.error("Incoming queue full, message dropped", message_id=message.id)
                    except InvalidJsonRpcError as e:
                        logger.error("Invalid JSON-RPC message", error=str(e))
                        # Continue reading despite parse errors
                        
                except asyncio.CancelledError:
                    logger.info("Read loop cancelled")
                    break
                except Exception as e:
                    logger.error("Error in read loop", error=str(e), exc_info=True)
                    # Continue reading despite errors
                    await asyncio.sleep(0.1)
        
        finally:
            logger.info("Read loop stopped")

    async def _write_loop(self) -> None:
        """Background task for writing messages to stdout.
        
        Continuously processes the outgoing message queue and writes
        messages to stdout in newline-delimited format. Messages are
        sent in FIFO order to maintain message ordering.
        """
        logger.info("Write loop started")
        
        try:
            while self._running:
                try:
                    # Wait for a message with timeout to allow checking shutdown
                    try:
                        message = await asyncio.wait_for(
                            self._outgoing_queue.get(),
                            timeout=1.0
                        )
                    except asyncio.TimeoutError:
                        continue
                    
                    if not self._stdout_writer:
                        logger.error("Stdout writer not available")
                        continue
                    
                    # Serialize and write the message
                    line = self._serializer.serialize(message)
                    self._stdout_writer.write(line.encode('utf-8'))
                    await self._stdout_writer.drain()
                    
                    logger.debug("Wrote message to stdout", message_id=message.id)
                    
                except asyncio.CancelledError:
                    logger.info("Write loop cancelled")
                    break
                except Exception as e:
                    logger.error("Error in write loop", error=str(e), exc_info=True)
                    await asyncio.sleep(0.1)
        
        finally:
            logger.info("Write loop stopped")

    async def _process_loop(self) -> None:
        """Background task for processing incoming messages.
        
        Continuously processes incoming messages using the message handler
        and queues responses for sending. Handles both requests (which expect
        responses) and notifications (which don't).
        
        Messages are processed in FIFO order to maintain message ordering
        as required by the MCP protocol.
        """
        logger.info("Process loop started")
        
        if not self._message_handler:
            logger.warning("No message handler set, messages will not be processed")
            # Still need to consume messages to prevent queue from filling up
            while self._running:
                try:
                    try:
                        message = await asyncio.wait_for(
                            self._incoming_queue.get(),
                            timeout=1.0
                        )
                        logger.warning("Message received but no handler set", message_id=message.id)
                    except asyncio.TimeoutError:
                        continue
                except asyncio.CancelledError:
                    break
            return
        
        try:
            while self._running:
                try:
                    # Wait for a message with timeout to allow checking shutdown
                    try:
                        message = await asyncio.wait_for(
                            self._incoming_queue.get(),
                            timeout=1.0
                        )
                    except asyncio.TimeoutError:
                        continue
                    
                    # Process the message using the handler
                    try:
                        response = await self._message_handler(message)
                        
                        # Queue the response for sending (if any)
                        # Notifications (messages without ID) don't get responses
                        if response is not None:
                            try:
                                self._outgoing_queue.put_nowait(response)
                                logger.debug("Processed message, queued response", 
                                           request_id=message.id, response_id=response.id)
                            except asyncio.QueueFull:
                                logger.error("Outgoing queue full, response dropped",
                                           message_id=message.id)
                        elif message.is_notification():
                            logger.debug("Processed notification", method=message.method)
                        else:
                            logger.debug("Handler returned no response", message_id=message.id)
                            
                    except Exception as e:
                        logger.error("Error processing message", 
                                   message_id=message.id, error=str(e), exc_info=True)
                        # Create error response for requests (not notifications)
                        if message.id is not None:
                            from .jsonrpc import create_error_response
                            error_response = create_error_response(
                                request_id=message.id,
                                code=-32603,
                                message=f"Internal error: {str(e)}"
                            )
                            try:
                                self._outgoing_queue.put_nowait(error_response)
                            except asyncio.QueueFull:
                                logger.error("Outgoing queue full, error response dropped",
                                           message_id=message.id)
                    
                except asyncio.CancelledError:
                    logger.info("Process loop cancelled")
                    break
                except Exception as e:
                    logger.error("Error in process loop", error=str(e), exc_info=True)
                    await asyncio.sleep(0.1)
        
        finally:
            logger.info("Process loop stopped")

    # ========================================================================
    # Private Methods - Utilities
    # ========================================================================

    def _clear_queues(self) -> None:
        """Clear all message queues."""
        # Clear incoming queue
        while not self._incoming_queue.empty():
            try:
                self._incoming_queue.get_nowait()
            except asyncio.QueueEmpty:
                break
        
        # Clear outgoing queue
        while not self._outgoing_queue.empty():
            try:
                self._outgoing_queue.get_nowait()
            except asyncio.QueueEmpty:
                break
        
        logger.debug("Message queues cleared")
