#!/usr/bin/env python3
"""
Interactive Shell Function

Provides a generalized interactive shell interface using pexpect for
managing long-running shell sessions and executing commands with
real-time output handling.
"""

import logging
import os
import re
import time
from typing import Dict, Any, Optional, List, Union
import pexpect

logger = logging.getLogger(__name__)


class InteractiveShell:
    """Manages an interactive shell session using pexpect"""

    def __init__(
        self,
        shell_cmd: str = "bash",
        shell_args: List[str] = None,
        cwd: str = None,
        timeout: int = 300,
        prompt_pattern: str = None,
        custom_prompt: str = None,
        encoding: str = "utf-8",
    ):
        """
        Initialize an interactive shell session

        Args:
            shell_cmd: Command to start the shell (default: "bash")
            shell_args: Arguments for the shell command
            cwd: Working directory for the shell
            timeout: Default timeout for operations in seconds
            prompt_pattern: Regex pattern to match shell prompts
            custom_prompt: Custom prompt to set (helps avoid ANSI sequences)
            encoding: Text encoding for the shell
        """
        self.shell_cmd = shell_cmd
        self.shell_args = shell_args or ["--norc", "--noprofile"]
        self.cwd = cwd or os.getcwd()
        self.timeout = timeout
        self.prompt_pattern = prompt_pattern or [r"\$", r"#", r">", r"nsh>", r"px4>"]
        self.custom_prompt = custom_prompt
        self.encoding = encoding

        self.process = None
        self.is_active = False
        self.created_at = time.time()
        self.command_history = []

    def start(self) -> Dict[str, Any]:
        """Start the interactive shell session"""
        try:
            # Start the shell process
            self.process = pexpect.spawn(
                self.shell_cmd,
                self.shell_args,
                cwd=self.cwd,
                timeout=self.timeout,
                encoding=self.encoding,
            )

            # Wait for initial shell prompt
            initial_timeout = 10
            try:
                self.process.expect(self.prompt_pattern, timeout=initial_timeout)
            except pexpect.TIMEOUT:
                output = self.process.before or ""
                return {
                    "success": False,
                    "error": f"Timeout waiting for shell prompt after {initial_timeout}s",
                    "output": output,
                }

            # Set custom prompt if specified
            if self.custom_prompt:
                self._set_custom_prompt()

            # Clear any remaining output
            self._clear_buffer()

            self.is_active = True
            logger.info(f"Started interactive shell: {self.shell_cmd} in {self.cwd}")

            return {
                "success": True,
                "shell_cmd": self.shell_cmd,
                "cwd": self.cwd,
                "message": "Interactive shell started successfully",
            }

        except Exception as e:
            logger.error(f"Failed to start shell: {e}")
            self.is_active = False
            return {
                "success": False,
                "error": f"Shell startup failed: {str(e)}",
                "shell_cmd": self.shell_cmd,
                "shell_args": self.shell_args,
                "cwd": self.cwd,
            }

    def execute_command(
        self,
        command: str,
        wait: int = 0,
        expect_patterns: Optional[List[str]] = None,
        responses: Optional[List[str]] = None,
        capture_output: bool = True,
        continue_on_error: bool = False,
        require_all_patterns: bool = True,
    ) -> Dict[str, Any]:
        """
        Execute a command in the interactive shell

        Args:
            command: Command to execute
            wait: Wait time in seconds. If 0, complete as soon as prompt appears.
                  If > 0, always wait this long and return output (ignore prompts)
            expect_patterns: Additional patterns to expect (beyond prompt)
            responses: Responses to send when patterns are matched (must match length of expect_patterns if provided)
            capture_output: Whether to capture and return output
            continue_on_error: Whether to continue if command fails
            require_all_patterns: If True (default), all expect_patterns must be matched for success.
                                  If False, patterns are optional - command succeeds even if patterns don't appear.

        Returns:
            Dictionary with execution results
        """
        if not self.is_active or not self.process:
            return {
                "success": False,
                "error": "Shell session is not active",
                "command": command,
            }

        command_start_time = time.time()

        try:
            # Validate responses match expect_patterns if both provided
            if responses is not None and expect_patterns is not None:
                if len(responses) != len(expect_patterns):
                    return {
                        "success": False,
                        "error": f"Number of responses ({len(responses)}) must match number of expect_patterns ({len(expect_patterns)})",
                        "command": command,
                    }

            # Clear buffer before sending command
            self._clear_buffer()

            # Send the command
            logger.info(f"Executing command: {command}")
            self.process.sendline(command)

            # Wait for the command echo to be consumed before looking for prompts
            # This prevents matching prompt patterns (like $) that appear in the command itself
            self._consume_command_echo(command)

            # If we have expect patterns with responses, use interactive response handling
            if expect_patterns and responses:
                return self._execute_with_interactive_responses(
                    command, expect_patterns, responses, capture_output, command_start_time, require_all_patterns
                )

            # If wait > 0, use simple time-based approach
            if wait > 0:
                return self._execute_with_fixed_wait(
                    command, wait, capture_output, command_start_time
                )

            # Determine what patterns to expect
            expect_patterns = expect_patterns or []
            all_patterns = (
                self.prompt_pattern + expect_patterns + [pexpect.TIMEOUT, pexpect.EOF]
            )

            # For wait=0, wait for prompt to appear
            return self._execute_with_prompt_wait(
                command, expect_patterns, capture_output, command_start_time
            )

        except Exception as e:
            logger.error(f"Error executing command '{command}': {e}")
            duration = time.time() - command_start_time
            return {
                "success": False,
                "error": f"Command execution failed: {str(e)}",
                "command": command,
                "duration": duration,
                "session_active": self.is_active,
                "process_alive": self.process.isalive() if self.process else False,
            }

    def _execute_with_fixed_wait(
        self, command: str, wait: int, capture_output: bool, command_start_time: float
    ) -> Dict[str, Any]:
        """Execute command and wait for fixed time, ignoring prompts"""
        logger.info(f"Waiting {wait} seconds for command to complete")

        all_output = ""

        # Collect output for the specified wait time
        end_time = command_start_time + wait
        while time.time() < end_time:
            try:
                remaining = end_time - time.time()
                chunk = self.process.read_nonblocking(
                    size=1000, timeout=min(remaining, 1)
                )
                if capture_output and chunk:
                    all_output += chunk
            except pexpect.TIMEOUT:
                continue  # Keep waiting
            except pexpect.EOF:
                break

        # Clean up output
        clean_output = self._clean_output(all_output, command) if capture_output else ""
        duration = time.time() - command_start_time

        # Always consider fixed-wait commands successful
        result = {
            "success": True,
            "command": command,
            "output": clean_output,
            "duration": duration,
            "cwd": self.cwd,
        }

        self.command_history.append(result)
        return result

    def _execute_with_prompt_wait(
        self,
        command: str,
        expect_patterns: List[str],
        capture_output: bool,
        command_start_time: float,
    ) -> Dict[str, Any]:
        """Execute command and wait for prompt to appear"""
        all_patterns = (
            self.prompt_pattern
            + (expect_patterns or [])
            + [pexpect.TIMEOUT, pexpect.EOF]
        )
        all_output = ""
        command_completed = False

        # Use a reasonable default timeout for prompt detection
        max_wait = 60  # 60 seconds max

        while not command_completed:
            try:
                elapsed = time.time() - command_start_time
                if elapsed > max_wait:
                    logger.warning(f"Command '{command}' exceeded maximum wait time")
                    break

                # Wait for patterns
                index = self.process.expect(all_patterns, timeout=5)

                # Collect output
                if capture_output and self.process.before:
                    all_output += self.process.before

                # Check what we matched
                if index < len(self.prompt_pattern):
                    # Found a prompt - command is done
                    # Trust pexpect's match since we consumed the command echo first
                    command_completed = True
                    logger.info(f"Command completed after {elapsed:.2f}s")

                elif index < len(self.prompt_pattern) + len(expect_patterns or []):
                    # Matched a custom pattern - continue reading
                    continue

                elif index >= len(all_patterns) - 2:  # TIMEOUT or EOF
                    # Continue unless we've exceeded max wait
                    continue

            except pexpect.TIMEOUT:
                continue
            except pexpect.EOF:
                if self.process.before and capture_output:
                    all_output += self.process.before
                break

        # Clean up output and determine success
        clean_output = self._clean_output(all_output, command) if capture_output else ""
        success = self._check_command_success(clean_output, command, False)
        duration = time.time() - command_start_time

        # Store result
        result = {
            "success": success,
            "command": command,
            "output": clean_output,
            "duration": duration,
            "cwd": self.cwd,
        }

        # Add error details if command failed
        if not success:
            error_info = self._extract_error_details(clean_output, command)
            result["error"] = error_info

        self.command_history.append(result)
        return result

    def _execute_with_interactive_responses(
        self,
        command: str,
        expect_patterns: List[str],
        responses: List[str],
        capture_output: bool,
        command_start_time: float,
        require_all_patterns: bool = True,
    ) -> Dict[str, Any]:
        """Execute command and send responses when patterns are matched"""
        all_output = ""
        max_wait = 120  # 2 minutes max for interactive commands

        patterns_matched = set()

        # Escape patterns for literal matching (users shouldn't need to write regex)
        escaped_patterns = [re.escape(p) for p in expect_patterns]

        # Ensure prompt_pattern is a list
        prompt_patterns = self.prompt_pattern if isinstance(self.prompt_pattern, list) else [self.prompt_pattern]

        # Build the list of all patterns to watch for (including prompts)
        all_patterns = (
            escaped_patterns
            + prompt_patterns
            + [pexpect.TIMEOUT, pexpect.EOF]
        )

        logger.info(f"Interactive command with {len(expect_patterns)} expected patterns")

        # Track authentication failures to prevent account lockout
        auth_failure_patterns = [
            "sorry, try again",
            "permission denied",
            "authentication failure",
            "incorrect password",
            "access denied",
            "login incorrect",
        ]
        auth_failure_count = 0
        max_auth_failures = 1  # Bail out after first auth failure to prevent lockout

        while True:
            try:
                elapsed = time.time() - command_start_time
                if elapsed > max_wait:
                    logger.warning(f"Command '{command}' exceeded maximum wait time")
                    break

                # Wait for any pattern
                index = self.process.expect(all_patterns, timeout=10)

                # Collect output
                if capture_output and self.process.before:
                    all_output += self.process.before

                    # Check for authentication failures to prevent account lockout
                    recent_output = self.process.before.lower()
                    for failure_pattern in auth_failure_patterns:
                        if failure_pattern in recent_output:
                            auth_failure_count += 1
                            logger.warning(f"Authentication failure detected: '{failure_pattern}' (count: {auth_failure_count})")
                            if auth_failure_count >= max_auth_failures:
                                logger.error(f"Too many authentication failures ({auth_failure_count}), aborting to prevent account lockout")
                                # Return early with failure
                                clean_output = self._clean_output(all_output, command)
                                return {
                                    "success": False,
                                    "command": command,
                                    "output": clean_output,
                                    "duration": time.time() - command_start_time,
                                    "cwd": self.cwd,
                                    "patterns_matched": list(patterns_matched),
                                    "expected_patterns": expect_patterns,
                                    "error": f"Authentication failed - aborting to prevent account lockout. Output: {clean_output[-500:] if len(clean_output) > 500 else clean_output}",
                                }
                            break

                # Check what we matched
                if index < len(expect_patterns):
                    # Matched one of our expected patterns
                    matched_pattern = expect_patterns[index]
                    response = responses[index]

                    logger.info(f"Matched pattern {index}: '{matched_pattern}', sending response")

                    # Send the response
                    self.process.sendline(response)
                    patterns_matched.add(matched_pattern)

                    # If we've matched all patterns, we might be done, but continue
                    # until we see a prompt
                    if len(patterns_matched) == len(expect_patterns):
                        logger.info("All patterns matched, waiting for final prompt")

                elif index < len(expect_patterns) + len(prompt_patterns):
                    # Found a prompt - command is done
                    # Trust the pexpect match for interactive commands - the prompt
                    # was matched, so we're at a shell prompt
                    logger.info(f"Command completed after {elapsed:.2f}s")
                    break

                elif index >= len(all_patterns) - 2:  # TIMEOUT or EOF
                    # Continue unless we've exceeded max wait
                    if elapsed > max_wait:
                        break
                    continue

            except pexpect.TIMEOUT:
                elapsed = time.time() - command_start_time
                if elapsed > max_wait:
                    break
                continue
            except pexpect.EOF:
                if self.process.before and capture_output:
                    all_output += self.process.before
                break

        # Clean up output and determine success
        clean_output = self._clean_output(all_output, command) if capture_output else ""

        # Check if all expected patterns were matched (or if we don't require them all)
        all_matched = len(patterns_matched) == len(expect_patterns)
        if require_all_patterns:
            success = all_matched and self._check_command_success(clean_output, command, False)
        else:
            # Patterns are optional - just check command success
            success = self._check_command_success(clean_output, command, False)

        duration = time.time() - command_start_time

        # Store result
        result = {
            "success": success,
            "command": command,
            "output": clean_output,
            "duration": duration,
            "cwd": self.cwd,
            "patterns_matched": list(patterns_matched),
            "expected_patterns": expect_patterns,
        }

        # Add error details if command failed
        if not success:
            if require_all_patterns and not all_matched:
                missing = set(expect_patterns) - patterns_matched
                result["error"] = f"Not all patterns matched. Missing: {missing}"
            else:
                error_info = self._extract_error_details(clean_output, command)
                result["error"] = error_info

        self.command_history.append(result)
        return result

    def send_input(self, text: str) -> Dict[str, Any]:
        """Send input to the interactive shell without expecting a prompt"""
        if not self.is_active or not self.process:
            return {"success": False, "error": "Shell session is not active"}

        try:
            self.process.send(text)
            return {"success": True, "message": f"Sent input: {repr(text)}"}
        except Exception as e:
            return {"success": False, "error": str(e)}

    def read_output(self, timeout: int = 5) -> Dict[str, Any]:
        """Read available output from the shell"""
        if not self.is_active or not self.process:
            return {"success": False, "error": "Shell session is not active"}

        try:
            output = ""
            try:
                while True:
                    chunk = self.process.read_nonblocking(size=1000, timeout=timeout)
                    if not chunk:
                        break
                    output += chunk
                    timeout = 0.1  # Reduce timeout for subsequent reads
            except pexpect.TIMEOUT:
                pass  # Normal - no more data available

            return {"success": True, "output": output}

        except Exception as e:
            return {"success": False, "error": str(e)}

    def close(self) -> Dict[str, Any]:
        """Close the interactive shell session and kill any child processes"""
        try:
            if self.process and self.is_active:
                # Send exit command to shell first
                try:
                    self.process.sendline("exit")
                    self.process.expect(pexpect.EOF, timeout=5)
                except (pexpect.TIMEOUT, pexpect.EOF):
                    pass  # Shell may have already exited

                # Terminate the process and its children
                if self.process.isalive():
                    self.process.terminate()
                    try:
                        self.process.wait(timeout=5)
                    except pexpect.TIMEOUT:
                        # Force kill if it doesn't terminate gracefully
                        self.process.kill()

                # Kill any remaining child processes (like gazebo)
                self._kill_child_processes()

                self.process.close()
                self.is_active = False
                logger.info(
                    "Closed interactive shell session and cleaned up child processes"
                )
                return {
                    "success": True,
                    "message": "Shell session closed and child processes terminated",
                }
            return {"success": True, "message": "Shell session was already closed"}
        except Exception as e:
            logger.error(f"Error closing shell session: {e}")
            return {"success": False, "error": str(e)}

    def _kill_child_processes(self):
        """Kill child processes spawned by the shell"""
        try:
            import psutil
            import signal

            if not self.process or not hasattr(self.process, "pid"):
                return

            try:
                parent = psutil.Process(self.process.pid)
                children = parent.children(recursive=True)

                # First try to terminate children gracefully
                for child in children:
                    try:
                        child.terminate()
                        logger.info(
                            f"Terminated child process: {child.pid} ({child.name()})"
                        )
                    except (psutil.NoSuchProcess, psutil.AccessDenied):
                        pass

                # Wait a bit for graceful termination
                import time

                time.sleep(2)

                # Force kill any remaining children
                for child in children:
                    try:
                        if child.is_running():
                            child.kill()
                            logger.info(
                                f"Force killed child process: {child.pid} ({child.name()})"
                            )
                    except (psutil.NoSuchProcess, psutil.AccessDenied):
                        pass

            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass

        except ImportError:
            # psutil not available, try basic approach
            logger.warning("psutil not available, using basic process cleanup")
            try:
                import os
                import signal

                if self.process and hasattr(self.process, "pid"):
                    # Try to kill the process group
                    os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
                    time.sleep(2)
                    os.killpg(os.getpgid(self.process.pid), signal.SIGKILL)
            except Exception as e:
                logger.debug(f"Basic process cleanup failed: {e}")
        except Exception as e:
            logger.debug(f"Child process cleanup failed: {e}")

    def _extract_error_details(self, output: str, command: str) -> str:
        """Extract detailed error information from command output, focusing on the tail"""
        if not output:
            return f"Command '{command}' failed with no output"

        lines = output.split("\n")
        non_empty_lines = [line.strip() for line in lines if line.strip()]

        if not non_empty_lines:
            return f"Command '{command}' failed with no readable output"

        # Always return the last 10 lines of output for context
        # This is most useful for build errors, compilation failures, etc.
        tail_lines = non_empty_lines[-10:]

        # Also look for critical error patterns in the tail
        error_patterns = [
            "error:",
            "failed",
            "cannot",
            "permission denied",
            "no such file or directory",
            "command not found",
            "syntax error",
            "make: ***",
            "fatal:",
            "abort:",
            "exception:",
            "compilation terminated",
            "build failed",
        ]

        # Find error lines in the tail
        error_lines = []
        for line in tail_lines:
            line_lower = line.lower()
            if any(pattern in line_lower for pattern in error_patterns):
                error_lines.append(line)

        # If we found specific errors in tail, prioritize those
        if error_lines:
            return f"Errors found: {' | '.join(error_lines[:5])}"

        # Otherwise return the complete tail for context
        return f"Command failed - Last 10 lines: {' | '.join(tail_lines)}"

    def _cleanup_on_error(self):
        """Clean up processes when command fails or times out"""
        logger.info("Cleaning up processes due to error/timeout")
        try:
            # Kill child processes but keep the shell session alive if possible
            self._kill_child_processes()
        except Exception as e:
            logger.debug(f"Error during cleanup: {e}")

    def get_status(self) -> Dict[str, Any]:
        """Get current status of the shell session"""
        return {
            "is_active": self.is_active,
            "shell_cmd": self.shell_cmd,
            "cwd": self.cwd,
            "created_at": self.created_at,
            "commands_executed": len(self.command_history),
            "process_alive": self.process.isalive() if self.process else False,
        }

    def _set_custom_prompt(self):
        """Set a custom prompt to avoid ANSI escape sequences"""
        try:
            if self.custom_prompt:
                self.process.sendline(f'export PS1="{self.custom_prompt}"')
                # Add custom prompt to existing patterns (don't replace them)
                # This allows detecting both the custom prompt and other common prompts (useful for nested shells like SSH)
                custom_pattern = re.escape(self.custom_prompt)

                # Ensure prompt_pattern is a list
                if isinstance(self.prompt_pattern, str):
                    self.prompt_pattern = [self.prompt_pattern]

                if custom_pattern not in self.prompt_pattern:
                    self.prompt_pattern.insert(0, custom_pattern)  # Add at beginning for priority matching

                # Wait for new prompt
                self.process.expect([custom_pattern], timeout=10)
                # Clear any remaining output
                self.process.sendline("")
                self.process.expect([custom_pattern], timeout=5)
                logger.info(f"Set custom prompt: {self.custom_prompt}")
        except Exception as e:
            logger.warning(f"Failed to set custom prompt: {e}")

    def _clear_buffer(self):
        """Clear any accumulated output in the process buffer"""
        if not self.process or not self.is_active:
            return

        try:
            while True:
                try:
                    discarded = self.process.read_nonblocking(size=1000, timeout=0.1)
                    if not discarded:
                        break
                    logger.debug(
                        f"Discarded buffer content: {repr(discarded[:100])}..."
                    )
                except pexpect.TIMEOUT:
                    break
                except pexpect.EOF:
                    break
        except Exception as e:
            logger.debug(f"Error clearing buffer: {e}")

    def _consume_command_echo(self, command: str, timeout: float = 2.0):
        """
        Wait for the command echo to be consumed before looking for prompts.

        This prevents matching prompt patterns (like $) that appear in the command itself.
        For example, 'sudo kill -9 $(pgrep foo)' contains a $ that could match prompt patterns.
        """
        if not self.process or not self.is_active:
            return

        try:
            # Read output until we see the command (or a significant portion of it)
            # or until timeout. We just need to get past the echo.
            collected = ""
            start_time = time.time()

            # Use the first distinctive part of the command for matching
            # (commands may be wrapped or truncated in the echo)
            cmd_start = command[:min(20, len(command))]

            while time.time() - start_time < timeout:
                try:
                    chunk = self.process.read_nonblocking(size=100, timeout=0.1)
                    if chunk:
                        collected += chunk
                        # Check if we've seen the command echo
                        if cmd_start in collected or command in collected:
                            logger.debug(f"Command echo consumed: {repr(collected[:100])}...")
                            return
                except pexpect.TIMEOUT:
                    # Small timeout, keep trying
                    continue
                except pexpect.EOF:
                    break

            logger.debug(f"Command echo timeout, collected: {repr(collected[:100])}...")
        except Exception as e:
            logger.debug(f"Error consuming command echo: {e}")

    def _clean_output(self, raw_output: str, command: str) -> str:
        """Clean up raw pexpect output"""
        if not raw_output:
            return ""

        # Remove ANSI escape sequences
        cleaned = re.sub(r"\x1b\[[?]?[0-9;]*[a-zA-Z]", "", raw_output)

        # Remove carriage returns
        cleaned = cleaned.replace("\r", "")

        # For commands that produce important status output, preserve more content
        # This can be configured by caller through command patterns or other means

        # For other commands, use the original cleaning logic
        lines = cleaned.split("\n")
        clean_lines = []

        skip_command_echo = False

        for line in lines:
            # Skip the echoed command line
            if command in line and not skip_command_echo:
                skip_command_echo = True
                continue
            elif skip_command_echo:
                clean_lines.append(line)

        # If we didn't find the command echo, return all lines
        if not skip_command_echo:
            clean_lines = lines

        # Remove empty lines at start and end
        while clean_lines and not clean_lines[0].strip():
            clean_lines.pop(0)
        while clean_lines and not clean_lines[-1].strip():
            clean_lines.pop()

        return "\n".join(clean_lines)

    def _check_command_success(
        self, output: str, command: str, continue_on_error: bool = False
    ) -> bool:
        """Check if command was successful based on output"""
        if continue_on_error:
            return True

        if not output:
            return True  # No output might be success for some commands

        output_lower = output.lower()

        # Common error indicators
        error_indicators = [
            "error:",
            "failed",
            "cannot",
            "permission denied",
            "no such file or directory",
            "command not found",
            "syntax error",
            "make: *** [",
            "compilation terminated",
            "build failed",
            "fatal error",
            "abort",
            "segmentation fault",
        ]

        for indicator in error_indicators:
            if indicator in output_lower:
                return False

        return True


class ShellSessionManager:
    """Manages multiple interactive shell sessions"""

    def __init__(self):
        self.sessions: Dict[str, InteractiveShell] = {}

    def create_session(
        self, session_id: str, shell_cmd: str = "bash", **kwargs
    ) -> Dict[str, Any]:
        """Create a new shell session"""
        if session_id in self.sessions:
            return {"success": False, "error": f"Session '{session_id}' already exists"}

        try:
            session = InteractiveShell(shell_cmd=shell_cmd, **kwargs)
            result = session.start()

            if result["success"]:
                self.sessions[session_id] = session
                result["session_id"] = session_id
            else:
                result["error"] = (
                    f"Failed to create session '{session_id}': {result.get('error', 'Unknown error')}"
                )

            return result
        except Exception as e:
            return {
                "success": False,
                "error": f"Exception creating session '{session_id}': {str(e)}",
                "shell_cmd": shell_cmd,
                "kwargs": kwargs,
            }

    def get_session(self, session_id: str) -> Optional[InteractiveShell]:
        """Get an existing shell session"""
        return self.sessions.get(session_id)

    def execute_in_session(
        self, session_id: str, command: str, **kwargs
    ) -> Dict[str, Any]:
        """Execute command in a specific session"""
        session = self.get_session(session_id)
        if not session:
            return {
                "success": False,
                "error": f"Session '{session_id}' not found in active sessions: {list(self.sessions.keys())}",
            }

        try:
            result = session.execute_command(command, **kwargs)
            result["session_id"] = session_id
            return result
        except Exception as e:
            return {
                "success": False,
                "error": f"Exception executing command in session '{session_id}': {str(e)}",
                "command": command,
                "session_id": session_id,
            }

    def close_session(self, session_id: str) -> Dict[str, Any]:
        """Close and remove a session"""
        session = self.sessions.pop(session_id, None)
        if not session:
            return {"success": False, "error": f"Session '{session_id}' not found"}

        result = session.close()
        result["session_id"] = session_id
        return result

    def list_sessions(self) -> Dict[str, Any]:
        """List all active sessions"""
        sessions_info = {}
        for session_id, session in self.sessions.items():
            sessions_info[session_id] = session.get_status()

        return {"success": True, "sessions": sessions_info}

    def cleanup_inactive_sessions(self) -> Dict[str, Any]:
        """Remove inactive sessions"""
        inactive_sessions = []
        for session_id, session in list(self.sessions.items()):
            if not session.is_active or (
                session.process and not session.process.isalive()
            ):
                inactive_sessions.append(session_id)
                session.close()
                del self.sessions[session_id]

        return {
            "success": True,
            "cleaned_sessions": inactive_sessions,
            "message": f"Cleaned up {len(inactive_sessions)} inactive sessions",
        }

    def cleanup_all_sessions(self) -> Dict[str, Any]:
        """Force close all sessions and clean up their processes"""
        closed_sessions = []
        for session_id, session in list(self.sessions.items()):
            try:
                session.close()
                closed_sessions.append(session_id)
            except Exception as e:
                logger.error(f"Error closing session {session_id}: {e}")

        self.sessions.clear()

        return {
            "success": True,
            "closed_sessions": closed_sessions,
            "message": f"Force closed {len(closed_sessions)} sessions",
        }


# Global session manager
session_manager = ShellSessionManager()
