#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Process Service Implementation using Connect RPC

Cross-platform support:
- Unix: Uses pty, fcntl, termios for PTY mode
- Windows: Uses pywinpty (ConPTY) via platform abstraction layer
"""

import asyncio
import os
import signal as sig
import subprocess
import sys
from typing import Dict, Optional

import psutil

# Platform abstraction for cross-platform PTY support
from executor.platform_compat import IS_WINDOWS, get_pty_manager, get_signal_handler
from shared.logger import setup_logger

# Will be generated by generate.sh
try:
    from ..gen.process.process import process_pb2
except ImportError:
    process_pb2 = None

logger = setup_logger("process_service")


class ConnectError(Exception):
    """Connect RPC error"""

    def __init__(self, code: str, message: str):
        self.code = code
        self.message = message
        super().__init__(message)


class ProcessManager:
    """Manages running processes"""

    def __init__(self):
        self.processes: Dict[int, subprocess.Popen] = {}
        self.tagged_processes: Dict[str, int] = {}  # tag -> pid mapping
        self.pty_fds: Dict[int, int] = {}  # pid -> master fd mapping for PTY processes

    def add_process(
        self,
        pid: int,
        process: subprocess.Popen,
        tag: Optional[str] = None,
        pty_fd: Optional[int] = None,
    ):
        """Add a process to the manager"""
        self.processes[pid] = process
        if tag:
            self.tagged_processes[tag] = pid
        if pty_fd is not None:
            self.pty_fds[pid] = pty_fd

    def get_pty_fd(self, selector: process_pb2.ProcessSelector) -> Optional[int]:
        """Get PTY master fd for a process"""
        pid = self.get_pid(selector)
        if pid:
            return self.pty_fds.get(pid)
        return None

    def get_process(
        self, selector: process_pb2.ProcessSelector
    ) -> Optional[subprocess.Popen]:
        """Get a process by selector (pid or tag)"""
        if selector.HasField("pid"):
            return self.processes.get(selector.pid)
        elif selector.HasField("tag"):
            pid = self.tagged_processes.get(selector.tag)
            if pid:
                return self.processes.get(pid)
        return None

    def get_pid(self, selector: process_pb2.ProcessSelector) -> Optional[int]:
        """Get PID from selector"""
        if selector.HasField("pid"):
            return selector.pid
        elif selector.HasField("tag"):
            return self.tagged_processes.get(selector.tag)
        return None

    def list_processes(self):
        """List all managed processes"""
        return [
            (pid, proc, tag if tag and self.tagged_processes.get(tag) == pid else None)
            for pid, proc in self.processes.items()
            for tag in [None] + list(self.tagged_processes.keys())
        ]

    def remove_process(self, pid: int):
        """Remove a process from the manager"""
        if pid in self.processes:
            del self.processes[pid]
        # Remove from tagged_processes
        tags_to_remove = [tag for tag, p in self.tagged_processes.items() if p == pid]
        for tag in tags_to_remove:
            del self.tagged_processes[tag]
        # Remove PTY fd if exists
        if pid in self.pty_fds:
            del self.pty_fds[pid]


class ProcessServiceHandler:
    """Handler for Process service RPC methods"""

    def __init__(self):
        self.manager = ProcessManager()

    async def List(self, request: process_pb2.ListRequest) -> process_pb2.ListResponse:
        """List all running processes in the system"""
        logger.debug("List request received")

        processes = []

        try:
            # Iterate through all system processes
            for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
                try:
                    pinfo = proc.info
                    pid = pinfo["pid"]
                    cmdline = pinfo.get("cmdline") or []

                    # Skip kernel processes with no cmdline
                    if not cmdline:
                        continue

                    # Build ProcessConfig
                    cmd = cmdline[0] if cmdline else ""
                    args = cmdline[1:] if len(cmdline) > 1 else []

                    # Check if this process was started by our manager
                    tag = None
                    for t, managed_pid in self.manager.tagged_processes.items():
                        if managed_pid == pid:
                            tag = t
                            break

                    info = process_pb2.ProcessInfo(
                        pid=pid,
                        config=process_pb2.ProcessConfig(
                            cmd=cmd,
                            args=args,
                        ),
                    )

                    if tag:
                        info.tag = tag

                    processes.append(info)

                except (
                    psutil.NoSuchProcess,
                    psutil.AccessDenied,
                    psutil.ZombieProcess,
                ):
                    # Skip processes that disappeared or we don't have access to
                    continue

        except Exception as e:
            logger.warning(f"Error listing processes: {e}")

        return process_pb2.ListResponse(processes=processes)

    async def Start(self, request: process_pb2.StartRequest):
        """Start a new process and stream its output"""
        logger.info(
            f"Starting process: {request.process.cmd} {list(request.process.args)}"
        )

        try:
            # Prepare environment
            env = os.environ.copy()
            if request.process.envs:
                env.update(request.process.envs)

            # Prepare command
            cmd = [request.process.cmd] + list(request.process.args)

            # Determine working directory
            cwd = request.process.cwd if request.process.cwd else None

            # Start process
            if request.HasField("pty"):
                # PTY mode - use platform abstraction for cross-platform support
                pty_manager = get_pty_manager()

                if not pty_manager.is_available():
                    raise ConnectError(
                        code="unavailable",
                        message="PTY mode is not available on this platform",
                    )

                # Spawn process with PTY using platform-specific implementation
                pty_process = pty_manager.spawn(
                    cmd=cmd,
                    cwd=cwd,
                    env=env,
                    rows=request.pty.size.rows if request.pty.HasField("size") else 24,
                    cols=request.pty.size.cols if request.pty.HasField("size") else 80,
                )

                # Get the underlying subprocess for process management
                # On Unix, we need to track the subprocess; on Windows, PtyProcess handles it
                if IS_WINDOWS:
                    # Windows: create a dummy Popen-like wrapper
                    proc = _WindowsPtyProcessWrapper(pty_process)
                    master = -1  # Not applicable on Windows
                else:
                    # Unix: get the actual process and fd
                    proc = pty_process._process
                    master = pty_process.fd

                # Store process with PTY master fd (or -1 on Windows)
                self.manager.add_process(
                    pty_process.pid,
                    proc,
                    request.tag if request.HasField("tag") else None,
                    pty_fd=master,
                )
                # Also store the pty_process for later use
                self.manager.pty_processes = getattr(self.manager, "pty_processes", {})
                self.manager.pty_processes[pty_process.pid] = pty_process

                # Send start event
                yield process_pb2.StartResponse(
                    event=process_pb2.ProcessEvent(
                        start=process_pb2.ProcessEvent.StartEvent(pid=proc.pid)
                    )
                )

                # Create tasks for reading PTY and sending keepalive
                loop = asyncio.get_event_loop()
                process_running = True
                last_data_time = loop.time()
                keepalive_interval = 5.0  # 5 seconds
                exit_signal = asyncio.Event()  # Signal for immediate exit

                async def read_pty_output():
                    nonlocal process_running, last_data_time

                    # Get the pty_process for reading
                    pty_proc = self.manager.pty_processes.get(pty_process.pid)
                    if not pty_proc:
                        return

                    if IS_WINDOWS:
                        # Windows: use PtyProcess.read() directly
                        try:
                            while pty_proc.poll() is None:
                                try:
                                    # Read with small timeout to allow checking process state
                                    data = pty_proc.read(4096)
                                    if data:
                                        last_data_time = loop.time()
                                        yield process_pb2.StartResponse(
                                            event=process_pb2.ProcessEvent(
                                                data=process_pb2.ProcessEvent.DataEvent(
                                                    pty=data
                                                )
                                            )
                                        )
                                    else:
                                        await asyncio.sleep(0.1)
                                except Exception:
                                    await asyncio.sleep(0.1)
                        finally:
                            process_running = False
                            exit_signal.set()
                    else:
                        # Unix: use select for non-blocking I/O
                        import fcntl
                        import select

                        # Set master fd to non-blocking
                        flags = fcntl.fcntl(master, fcntl.F_GETFL)
                        fcntl.fcntl(master, fcntl.F_SETFL, flags | os.O_NONBLOCK)

                        try:
                            while proc.poll() is None:
                                try:
                                    # Use select to check if data is available with timeout
                                    ready, _, _ = await loop.run_in_executor(
                                        None, select.select, [master], [], [], 0.5
                                    )
                                    if ready:
                                        try:
                                            data = os.read(master, 4096)
                                            if data:
                                                last_data_time = loop.time()
                                                yield process_pb2.StartResponse(
                                                    event=process_pb2.ProcessEvent(
                                                        data=process_pb2.ProcessEvent.DataEvent(
                                                            pty=data
                                                        )
                                                    )
                                                )
                                        except (OSError, BlockingIOError):
                                            pass
                                    else:
                                        await asyncio.sleep(0.1)
                                except (OSError, ValueError):
                                    break
                        finally:
                            process_running = False
                            exit_signal.set()

                async def send_keepalive():
                    keepalive_count = 0
                    while process_running and not exit_signal.is_set():
                        try:
                            # Wait for exit signal or timeout
                            await asyncio.wait_for(
                                exit_signal.wait(), timeout=keepalive_interval
                            )
                            # Signal received, exit immediately
                            break
                        except asyncio.TimeoutError:
                            # Normal timeout, check if keepalive needed
                            if (
                                process_running
                                and loop.time() - last_data_time >= keepalive_interval
                            ):
                                keepalive_count += 1
                                if keepalive_count % 6 == 1:  # Log every 30s (6 * 5s)
                                    logger.debug(
                                        f"PTY process {proc.pid} keepalive #{keepalive_count}, "
                                        f"no data for {loop.time() - last_data_time:.1f}s"
                                    )
                                yield process_pb2.StartResponse(
                                    event=process_pb2.ProcessEvent(
                                        keepalive=process_pb2.ProcessEvent.KeepAlive()
                                    )
                                )

                # Merge streams from PTY output and keepalive
                try:
                    async for response in self._merge_async_generators(
                        [read_pty_output(), send_keepalive()]
                    ):
                        yield response
                except Exception as e:
                    logger.exception(f"Error in PTY stream: {e}")
                # Note: Don't close master fd here - it's needed for SendInput/Update
                # It will be closed in cleanup when process exits

            else:
                # Normal mode - separate stdout/stderr
                proc = subprocess.Popen(
                    cmd,
                    stdin=subprocess.PIPE if request.stdin else subprocess.DEVNULL,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    env=env,
                    cwd=cwd,
                )

                # Store process
                self.manager.add_process(
                    proc.pid, proc, request.tag if request.HasField("tag") else None
                )

                # Send start event
                yield process_pb2.StartResponse(
                    event=process_pb2.ProcessEvent(
                        start=process_pb2.ProcessEvent.StartEvent(pid=proc.pid)
                    )
                )

                # Track data events for keepalive timing
                loop = asyncio.get_event_loop()
                last_data_time = loop.time()
                keepalive_interval = 5.0
                process_running = True
                exit_signal = asyncio.Event()  # Signal for immediate exit
                stderr_output = []  # Collect stderr for error logging

                # Stream output
                async def read_stream(stream, is_stderr=False):
                    nonlocal last_data_time
                    while True:
                        try:
                            data = await loop.run_in_executor(None, stream.read, 4096)
                            if not data:
                                break
                            last_data_time = loop.time()
                            if is_stderr:
                                stderr_output.append(
                                    data.decode("utf-8", errors="replace")
                                )
                                yield process_pb2.ProcessEvent.DataEvent(stderr=data)
                            else:
                                yield process_pb2.ProcessEvent.DataEvent(stdout=data)
                        except Exception as e:
                            logger.warning(f"Error reading stream: {e}")
                            break

                async def send_keepalive():
                    nonlocal process_running
                    keepalive_count = 0
                    while (
                        process_running
                        and proc.poll() is None
                        and not exit_signal.is_set()
                    ):
                        try:
                            # Wait for exit signal or timeout
                            await asyncio.wait_for(
                                exit_signal.wait(), timeout=keepalive_interval
                            )
                            # Signal received, exit immediately
                            break
                        except asyncio.TimeoutError:
                            # Normal timeout, check if keepalive needed
                            if (
                                process_running
                                and proc.poll() is None
                                and loop.time() - last_data_time >= keepalive_interval
                            ):
                                keepalive_count += 1
                                if keepalive_count % 6 == 1:  # Log every 30s (6 * 5s)
                                    logger.debug(
                                        f"Process {proc.pid} keepalive #{keepalive_count}, "
                                        f"no data for {loop.time() - last_data_time:.1f}s"
                                    )
                                yield process_pb2.StartResponse(
                                    event=process_pb2.ProcessEvent(
                                        keepalive=process_pb2.ProcessEvent.KeepAlive()
                                    )
                                )

                # Read both streams concurrently
                tasks = []
                if proc.stdout:
                    tasks.append(read_stream(proc.stdout, False))
                if proc.stderr:
                    tasks.append(read_stream(proc.stderr, True))

                async def data_stream():
                    try:
                        if tasks:
                            async for data_event in self._merge_streams(tasks):
                                yield process_pb2.StartResponse(
                                    event=process_pb2.ProcessEvent(data=data_event)
                                )
                    finally:
                        exit_signal.set()

                # Merge data stream and keepalive
                try:
                    async for response in self._merge_async_generators(
                        [data_stream(), send_keepalive()]
                    ):
                        yield response
                except Exception as e:
                    logger.exception(f"Error in stream merge: {e}")
                finally:
                    process_running = False
                    exit_signal.set()  # Signal keepalive to exit immediately

            # Wait for process to complete (non-blocking check)
            returncode = proc.poll()
            if returncode is None:
                # Process still running, wait asynchronously
                returncode = await loop.run_in_executor(None, proc.wait)

            if returncode != 0:
                stderr_text = "".join(stderr_output).strip()
                if stderr_text:
                    logger.warning(
                        f"Process {proc.pid} completed with non-zero exit code {returncode}, "
                        f"stderr: {stderr_text[:1000]}"
                    )
                else:
                    logger.warning(
                        f"Process {proc.pid} completed with non-zero exit code {returncode}"
                    )
            else:
                logger.info(f"Process {proc.pid} completed with exit code {returncode}")

            # Send end event
            yield process_pb2.StartResponse(
                event=process_pb2.ProcessEvent(
                    end=process_pb2.ProcessEvent.EndEvent(
                        exit_code=returncode,
                        exited=True,
                        status="completed",
                    )
                )
            )

            # Clean up - process will be removed by background cleanup
            # The manager will handle cleanup when process exits

        except Exception as e:
            logger.exception(f"Error starting process: {e}")
            yield process_pb2.StartResponse(
                event=process_pb2.ProcessEvent(
                    end=process_pb2.ProcessEvent.EndEvent(
                        exit_code=-1,
                        exited=False,
                        status="failed",
                        error=str(e),
                    )
                )
            )
        finally:
            # Background cleanup - only remove from manager if process has exited
            if "proc" in locals():
                asyncio.create_task(self._cleanup_finished_process(proc.pid))

    async def Connect(self, request: process_pb2.ConnectRequest):
        """Connect to an existing process and monitor its status"""
        logger.info(f"Connect request for process selector: {request.process}")

        # Try to get process from manager first (for processes started via Start)
        proc = self.manager.get_process(request.process)
        pid = self.manager.get_pid(request.process)

        # If not in manager, try to get from system by PID
        if not pid and request.process.HasField("pid"):
            pid = request.process.pid
            try:
                # Verify the process exists in the system
                psutil.Process(pid)
                proc = None  # System process, not managed by us
            except psutil.NoSuchProcess:
                raise ConnectError(
                    code="not_found",
                    message=f"Process with PID {pid} not found",
                )
        elif not pid:
            raise ConnectError(
                code="not_found",
                message="Process not found",
            )

        # Send start event with PID
        yield process_pb2.ConnectResponse(
            event=process_pb2.ProcessEvent(
                start=process_pb2.ProcessEvent.StartEvent(pid=pid)
            )
        )

        # Monitor process status
        try:
            if proc:
                # Managed process - use subprocess.Popen
                # Check if already finished
                if proc.poll() is not None:
                    logger.info(
                        f"Process {pid} already completed with exit code {proc.returncode}"
                    )
                    yield process_pb2.ConnectResponse(
                        event=process_pb2.ProcessEvent(
                            end=process_pb2.ProcessEvent.EndEvent(
                                exit_code=proc.returncode,
                                exited=True,
                                status="completed",
                            )
                        )
                    )
                    return

                # Send keepalive while running
                while proc.poll() is None:
                    yield process_pb2.ConnectResponse(
                        event=process_pb2.ProcessEvent(
                            keepalive=process_pb2.ProcessEvent.KeepAlive()
                        )
                    )
                    await asyncio.sleep(5)

                logger.info(f"Process {pid} completed with exit code {proc.returncode}")

                # Send end event
                yield process_pb2.ConnectResponse(
                    event=process_pb2.ProcessEvent(
                        end=process_pb2.ProcessEvent.EndEvent(
                            exit_code=proc.returncode,
                            exited=True,
                            status="completed",
                        )
                    )
                )
            else:
                # System process - use psutil
                try:
                    ps_proc = psutil.Process(pid)

                    # Monitor while process is running
                    while ps_proc.is_running():
                        yield process_pb2.ConnectResponse(
                            event=process_pb2.ProcessEvent(
                                keepalive=process_pb2.ProcessEvent.KeepAlive()
                            )
                        )
                        await asyncio.sleep(5)

                        # Refresh process info
                        try:
                            ps_proc.status()
                        except psutil.NoSuchProcess:
                            break

                    logger.info(f"System process {pid} ended")

                    # Process ended
                    yield process_pb2.ConnectResponse(
                        event=process_pb2.ProcessEvent(
                            end=process_pb2.ProcessEvent.EndEvent(
                                exit_code=-1,  # Unknown exit code for system processes
                                exited=True,
                                status="completed",
                            )
                        )
                    )
                except psutil.NoSuchProcess:
                    # Process already finished
                    yield process_pb2.ConnectResponse(
                        event=process_pb2.ProcessEvent(
                            end=process_pb2.ProcessEvent.EndEvent(
                                exit_code=-1,
                                exited=True,
                                status="completed",
                            )
                        )
                    )
        except Exception as e:
            logger.exception(f"Error in Connect: {e}")
            yield process_pb2.ConnectResponse(
                event=process_pb2.ProcessEvent(
                    end=process_pb2.ProcessEvent.EndEvent(
                        exit_code=-1,
                        exited=False,
                        status="failed",
                        error=str(e),
                    )
                )
            )

    async def Update(
        self, request: process_pb2.UpdateRequest
    ) -> process_pb2.UpdateResponse:
        """Update process settings (e.g., PTY size)"""
        logger.info(f"Update request for process selector: {request.process}")

        proc = self.manager.get_process(request.process)
        if not proc:
            raise ConnectError(
                code="not_found",
                message="Process not found",
            )

        # Update PTY size if provided
        if request.HasField("pty"):
            pty_fd = self.manager.get_pty_fd(request.process)
            pid = self.manager.get_pid(request.process)

            # Try to get pty_process for Windows support
            pty_processes = getattr(self.manager, "pty_processes", {})
            pty_proc = pty_processes.get(pid) if pid else None

            if pty_fd is None and pty_proc is None:
                raise ConnectError(
                    code="invalid_argument",
                    message="Process is not running in PTY mode",
                )

            try:
                rows = request.pty.size.rows
                cols = request.pty.size.cols

                if IS_WINDOWS and pty_proc:
                    # Windows: use PtyProcess.resize()
                    pty_proc.resize(rows, cols)
                elif pty_fd is not None:
                    # Unix: use ioctl
                    import fcntl
                    import struct
                    import termios

                    winsize = struct.pack("HHHH", rows, cols, 0, 0)
                    fcntl.ioctl(pty_fd, termios.TIOCSWINSZ, winsize)

                logger.debug(f"Resized PTY for PID {proc.pid} to {rows}x{cols}")
            except Exception as e:
                logger.exception(f"Error resizing PTY: {e}")
                raise ConnectError(
                    code="internal",
                    message=f"Failed to resize TTY: {str(e)}",
                )

        return process_pb2.UpdateResponse()

    def _handle_input(
        self,
        process_selector: process_pb2.ProcessSelector,
        input_data: process_pb2.ProcessInput,
    ):
        """Shared helper to handle input for both SendInput and StreamInput"""
        proc = self.manager.get_process(process_selector)
        if not proc:
            raise ConnectError(
                code="not_found",
                message="Process not found",
            )

        # Handle stdin input
        if input_data.HasField("stdin"):
            if not proc.stdin:
                raise ConnectError(
                    code="failed_precondition",
                    message="Process stdin not available",
                )
            try:
                proc.stdin.write(input_data.stdin)
                proc.stdin.flush()
            except Exception as e:
                logger.exception(f"Error writing to stdin: {e}")
                raise ConnectError(
                    code="internal",
                    message=f"Failed to write to stdin: {str(e)}",
                )

        # Handle PTY input
        elif input_data.HasField("pty"):
            pty_fd = self.manager.get_pty_fd(process_selector)
            if pty_fd is None:
                raise ConnectError(
                    code="failed_precondition",
                    message="Process is not running in PTY mode",
                )
            try:
                os.write(pty_fd, input_data.pty)
            except Exception as e:
                logger.exception(f"Error writing to PTY: {e}")
                raise ConnectError(
                    code="internal",
                    message=f"Failed to write to PTY: {str(e)}",
                )

    async def SendInput(
        self, request: process_pb2.SendInputRequest
    ) -> process_pb2.SendInputResponse:
        """Send input to a process"""
        logger.debug(f"SendInput request for process selector: {request.process}")

        self._handle_input(request.process, request.input)
        return process_pb2.SendInputResponse()

    async def StreamInput(self, request_iterator) -> process_pb2.StreamInputResponse:
        """Stream input to a process"""
        logger.info("StreamInput request received")

        process_selector = None
        async for request in request_iterator:
            if request.HasField("start"):
                # Initialize process connection
                process_selector = request.start.process
                # Verify process exists
                proc = self.manager.get_process(process_selector)
                if not proc:
                    raise ConnectError(
                        code="not_found",
                        message="Process not found",
                    )
            elif request.HasField("data"):
                if not process_selector:
                    raise ConnectError(
                        code="failed_precondition",
                        message="Process not started",
                    )
                # Handle input using shared helper
                self._handle_input(process_selector, request.data.input)
            elif request.HasField("keepalive"):
                # Just keep connection alive
                pass

        return process_pb2.StreamInputResponse()

    async def SendSignal(
        self, request: process_pb2.SendSignalRequest
    ) -> process_pb2.SendSignalResponse:
        """Send a signal to a process"""
        logger.info(
            f"SendSignal request for process selector: {request.process}, signal: {request.signal}"
        )

        # Get process
        proc = self.manager.get_process(request.process)
        if not proc:
            raise ConnectError(
                code="not_found",
                message="Process not found",
            )

        # Map proto signal to OS signal - use platform abstraction
        signal_handler = get_signal_handler()
        if request.signal == process_pb2.SIGNAL_SIGKILL:
            os_signal = signal_handler.get_kill_signal()
        elif request.signal == process_pb2.SIGNAL_SIGTERM:
            os_signal = signal_handler.get_termination_signal()
        else:
            raise ConnectError(
                code="invalid_argument",
                message=f"Invalid signal: {request.signal}",
            )

        # Send signal to process - use platform-specific method
        try:
            if IS_WINDOWS:
                # Windows: use platform signal handler
                if request.signal == process_pb2.SIGNAL_SIGKILL:
                    signal_handler.terminate_forcefully(proc.pid)
                else:
                    signal_handler.terminate_gracefully(proc.pid)
            else:
                # Unix: use subprocess.send_signal
                proc.send_signal(os_signal)
            logger.debug(f"Successfully sent signal {os_signal} to PID {proc.pid}")
        except Exception as e:
            logger.exception(f"Error sending signal to process: {e}")
            raise ConnectError(
                code="internal",
                message=f"Error sending signal: {str(e)}",
            )

        return process_pb2.SendSignalResponse()

    async def _merge_streams(self, generators):
        """Merge multiple async generators"""
        tasks = {}
        for gen in generators:
            try:
                task = asyncio.create_task(self._safe_anext(gen))
                tasks[task] = gen
            except Exception as e:
                logger.warning(f"Error creating task in merge_streams: {e}")

        while tasks:
            if not tasks:
                break

            done, _ = await asyncio.wait(
                tasks.keys(), return_when=asyncio.FIRST_COMPLETED
            )

            for task in done:
                gen = tasks.pop(task)
                try:
                    result = await task
                    if result is not None:
                        yield result
                        # Re-add task for next iteration
                        try:
                            new_task = asyncio.create_task(self._safe_anext(gen))
                            tasks[new_task] = gen
                        except Exception as e:
                            logger.warning(f"Error re-adding task: {e}")
                except StopAsyncIteration:
                    # Generator exhausted
                    pass
                except Exception as e:
                    # Log unexpected errors but continue
                    logger.warning(f"Error in merge_streams: {e}")

    async def _merge_async_generators(self, generators):
        """Merge multiple async generators that yield complete responses"""
        tasks = {}
        for gen in generators:
            try:
                task = asyncio.create_task(self._safe_anext(gen))
                tasks[task] = gen
            except Exception as e:
                logger.warning(f"Error creating task in merge_async_generators: {e}")

        while tasks:
            if not tasks:
                break

            done, _ = await asyncio.wait(
                tasks.keys(), return_when=asyncio.FIRST_COMPLETED
            )

            for task in done:
                gen = tasks.pop(task)
                try:
                    result = await task
                    if result is not None:
                        yield result
                        # Re-add task for next iteration
                        try:
                            new_task = asyncio.create_task(self._safe_anext(gen))
                            tasks[new_task] = gen
                        except Exception as e:
                            logger.warning(f"Error re-adding task: {e}")
                except StopAsyncIteration:
                    # Generator exhausted
                    pass
                except Exception as e:
                    # Log unexpected errors but continue
                    logger.warning(f"Error in merge_async_generators: {e}")

    async def _safe_anext(self, gen):
        """Safely get next item from async generator, catching StopAsyncIteration"""
        try:
            return await gen.__anext__()
        except StopAsyncIteration:
            # Signal that generator is exhausted by returning None
            return None

    async def _cleanup_process(self, pid: int):
        """Background cleanup of finished process"""
        await asyncio.sleep(1)  # Give process time to finish
        try:
            self.manager.remove_process(pid)
            logger.debug(f"Cleaned up process {pid}")
        except Exception as e:
            logger.warning(f"Error cleaning up process {pid}: {e}")

    async def _cleanup_finished_process(self, pid: int):
        """Cleanup process from manager only if it has exited"""
        try:
            # Check if process is in manager
            proc = None
            for p, process_obj in self.manager.processes.items():
                if p == pid:
                    proc = process_obj
                    break

            if not proc:
                # Not a managed process, nothing to clean up
                return

            # Wait for process to exit (with timeout)
            loop = asyncio.get_event_loop()
            max_wait = 300  # 5 minutes max
            start_time = loop.time()

            while proc.poll() is None:
                await asyncio.sleep(1)
                if loop.time() - start_time > max_wait:
                    logger.warning(
                        f"Process {pid} did not exit after {max_wait}s, removing from manager anyway"
                    )
                    break

            # Close PTY master fd if exists (Unix only)
            pty_fd = self.manager.pty_fds.get(pid)
            if pty_fd is not None and pty_fd != -1:
                try:
                    os.close(pty_fd)
                    logger.debug(f"Closed PTY fd {pty_fd} for process {pid}")
                except OSError as e:
                    logger.warning(f"Error closing PTY fd for process {pid}: {e}")

            # Close PtyProcess if exists (cross-platform)
            pty_processes = getattr(self.manager, "pty_processes", {})
            pty_proc = pty_processes.pop(pid, None)
            if pty_proc:
                try:
                    pty_proc.close()
                    logger.debug(f"Closed PtyProcess for process {pid}")
                except Exception as e:
                    logger.warning(f"Error closing PtyProcess for process {pid}: {e}")

            # Process has exited or timed out, remove from manager
            self.manager.remove_process(pid)
            exit_code = proc.returncode if proc else "unknown"
            logger.info(
                f"Cleaned up finished process pid: {pid} with exit code {exit_code}"
            )
        except Exception as e:
            logger.warning(f"Error cleaning up process {pid}: {e}")


class _WindowsPtyProcessWrapper:
    """Wrapper to make Windows PtyProcess compatible with subprocess.Popen interface.

    This allows the ProcessManager to work with Windows PTY processes
    using the same interface as Unix subprocess.Popen.
    """

    def __init__(self, pty_process):
        """Initialize wrapper with a PtyProcess instance.

        Args:
            pty_process: Platform PtyProcess instance from Windows PTY manager.
        """
        self._pty_process = pty_process
        self._returncode = None
        self.stdin = None  # PTY handles I/O differently

    @property
    def pid(self):
        """Get the process ID."""
        return self._pty_process.pid

    @property
    def returncode(self):
        """Get the return code, or None if process hasn't terminated."""
        if self._returncode is None:
            self._returncode = self._pty_process.poll()
        return self._returncode

    def poll(self):
        """Check if process has terminated.

        Returns:
            Exit code if terminated, None if still running.
        """
        result = self._pty_process.poll()
        if result is not None:
            self._returncode = result
        return result

    def wait(self, timeout=None):
        """Wait for process to terminate.

        Args:
            timeout: Maximum time to wait in seconds.

        Returns:
            Exit code of the process.
        """
        self._returncode = self._pty_process.wait(timeout=timeout)
        return self._returncode

    def terminate(self):
        """Terminate the process gracefully."""
        self._pty_process.terminate(force=False)

    def kill(self):
        """Forcefully kill the process."""
        self._pty_process.terminate(force=True)

    def send_signal(self, sig):
        """Send a signal to the process.

        On Windows, this translates to terminate/kill operations.

        Args:
            sig: Signal number (SIGTERM=15 for graceful, SIGKILL=9 for forceful).
        """
        if sig == 9:  # SIGKILL equivalent
            self._pty_process.terminate(force=True)
        else:
            self._pty_process.terminate(force=False)
