rp/pr/tools/process_handlers.py

from abc import ABC, abstractmethod


class ProcessHandler(ABC):
    """Base class for process-specific handlers."""

    def __init__(self, multiplexer):
        self.multiplexer = multiplexer
        self.state_machine = {}
        self.current_state = "initial"
        self.prompt_patterns = []
        self.response_suggestions = {}

    @abstractmethod
    def get_process_type(self):
        """Return the process type this handler manages."""
        pass

    def update_state(self, output):
        """Update internal state based on output."""
        pass

    def get_prompt_suggestions(self):
        """Return suggested responses for current state."""
        return self.response_suggestions.get(self.current_state, [])

    def is_waiting_for_input(self):
        """Check if process appears to be waiting for input."""
        return self.current_state in ["waiting_confirmation", "waiting_input"]


class AptHandler(ProcessHandler):
    """Handler for apt package manager interactions."""

    def __init__(self, multiplexer):
        super().__init__(multiplexer)
        self.state_machine = {
            "initial": ["running_command"],
            "running_command": ["waiting_confirmation", "completed"],
            "waiting_confirmation": ["confirmed", "cancelled"],
            "confirmed": ["installing", "completed"],
            "installing": ["completed", "error"],
            "completed": [],
            "error": [],
            "cancelled": [],
        }
        self.prompt_patterns = [
            (r"Do you want to continue\?", "confirmation"),
            (r"After this operation.*installed\.", "size_info"),
            (r"Need to get.*B of archives\.", "download_info"),
            (r"Unpacking.*Configuring", "configuring"),
            (r"Setting up", "setting_up"),
            (r"E:\s", "error"),
        ]

    def get_process_type(self):
        return "apt"

    def update_state(self, output):
        """Update state based on apt output patterns."""
        output_lower = output.lower()

        # Check for completion
        if "processing triggers" in output_lower or "done" in output_lower:
            self.current_state = "completed"
        # Check for confirmation prompts
        elif "do you want to continue" in output_lower:
            self.current_state = "waiting_confirmation"
        # Check for installation progress
        elif "setting up" in output_lower or "unpacking" in output_lower:
            self.current_state = "installing"
        # Check for errors
        elif "e:" in output_lower or "error" in output_lower:
            self.current_state = "error"

    def get_prompt_suggestions(self):
        """Return suggested responses for apt prompts."""
        suggestions = super().get_prompt_suggestions()
        if self.current_state == "waiting_confirmation":
            suggestions.extend(["y", "yes", "n", "no"])
        return suggestions


class VimHandler(ProcessHandler):
    """Handler for vim editor interactions."""

    def __init__(self, multiplexer):
        super().__init__(multiplexer)
        self.state_machine = {
            "initial": ["normal_mode", "insert_mode"],
            "normal_mode": ["insert_mode", "command_mode", "visual_mode"],
            "insert_mode": ["normal_mode"],
            "command_mode": ["normal_mode"],
            "visual_mode": ["normal_mode"],
            "exiting": [],
        }
        self.prompt_patterns = [
            (r"-- INSERT --", "insert_mode"),
            (r"-- VISUAL --", "visual_mode"),
            (r":", "command_mode"),
            (r"Press ENTER", "waiting_enter"),
            (r"Saved", "saved"),
        ]
        self.mode_indicators = {
            "insert": "-- INSERT --",
            "visual": "-- VISUAL --",
            "command": ":",
        }

    def get_process_type(self):
        return "vim"

    def update_state(self, output):
        """Update state based on vim mode indicators."""
        if "-- INSERT --" in output:
            self.current_state = "insert_mode"
        elif "-- VISUAL --" in output:
            self.current_state = "visual_mode"
        elif output.strip().endswith(":"):
            self.current_state = "command_mode"
        elif "Press ENTER" in output:
            self.current_state = "waiting_enter"
        else:
            # Default to normal mode if no specific indicators
            self.current_state = "normal_mode"

    def get_prompt_suggestions(self):
        """Return suggested commands for vim modes."""
        suggestions = super().get_prompt_suggestions()
        if self.current_state == "command_mode":
            suggestions.extend(["w", "q", "wq", "q!", "w!"])
        elif self.current_state == "normal_mode":
            suggestions.extend(["i", "a", "o", "dd", ":w", ":q"])
        elif self.current_state == "waiting_enter":
            suggestions.extend(["\n"])
        return suggestions


class SSHHandler(ProcessHandler):
    """Handler for SSH connection interactions."""

    def __init__(self, multiplexer):
        super().__init__(multiplexer)
        self.state_machine = {
            "initial": ["connecting"],
            "connecting": ["auth_prompt", "connected", "failed"],
            "auth_prompt": ["connected", "failed"],
            "connected": ["shell", "disconnected"],
            "shell": ["disconnected"],
            "failed": [],
            "disconnected": [],
        }
        self.prompt_patterns = [
            (r"password:", "password_prompt"),
            (r"yes/no", "host_key_prompt"),
            (r"Permission denied", "auth_failed"),
            (r"Welcome to", "connected"),
            (r"\$", "shell_prompt"),
            (r"\#", "root_shell_prompt"),
            (r"Connection closed", "disconnected"),
        ]

    def get_process_type(self):
        return "ssh"

    def update_state(self, output):
        """Update state based on SSH connection output."""
        output_lower = output.lower()

        if "permission denied" in output_lower:
            self.current_state = "failed"
        elif "password:" in output_lower:
            self.current_state = "auth_prompt"
        elif "yes/no" in output_lower:
            self.current_state = "auth_prompt"
        elif "welcome to" in output_lower or "last login" in output_lower:
            self.current_state = "connected"
        elif output.strip().endswith("$") or output.strip().endswith("#"):
            self.current_state = "shell"
        elif "connection closed" in output_lower:
            self.current_state = "disconnected"

    def get_prompt_suggestions(self):
        """Return suggested responses for SSH prompts."""
        suggestions = super().get_prompt_suggestions()
        if self.current_state == "auth_prompt":
            if "password:" in self.multiplexer.get_all_output()["stdout"]:
                suggestions.extend(["<password>"])  # Placeholder for actual password
            elif "yes/no" in self.multiplexer.get_all_output()["stdout"]:
                suggestions.extend(["yes", "no"])
        return suggestions


class GenericProcessHandler(ProcessHandler):
    """Fallback handler for unknown process types."""

    def __init__(self, multiplexer):
        super().__init__(multiplexer)
        self.state_machine = {
            "initial": ["running"],
            "running": ["waiting_input", "completed"],
            "waiting_input": ["running"],
            "completed": [],
        }
        self.prompt_patterns = [
            (r"\?\s*$", "waiting_input"),  # Lines ending with ?
            (r">\s*$", "waiting_input"),  # Lines ending with >
            (r":\s*$", "waiting_input"),  # Lines ending with :
            (r"done", "completed"),
            (r"finished", "completed"),
            (r"exit code", "completed"),
        ]

    def get_process_type(self):
        return "generic"

    def update_state(self, output):
        """Basic state detection for generic processes."""
        output_lower = output.lower()

        if any(pattern in output_lower for pattern in ["done", "finished", "complete"]):
            self.current_state = "completed"
        elif any(output.strip().endswith(char) for char in ["?", ">", ":"]):
            self.current_state = "waiting_input"
        else:
            self.current_state = "running"


# Handler registry
_handler_classes = {
    "apt": AptHandler,
    "vim": VimHandler,
    "ssh": SSHHandler,
    "generic": GenericProcessHandler,
}


def get_handler_for_process(process_type, multiplexer):
    """Get appropriate handler for a process type."""
    handler_class = _handler_classes.get(process_type, GenericProcessHandler)
    return handler_class(multiplexer)


def detect_process_type(command):
    """Detect process type from command."""
    command_str = " ".join(command) if isinstance(command, list) else command
    command_lower = command_str.lower()

    if "apt" in command_lower or "apt-get" in command_lower:
        return "apt"
    elif "vim" in command_lower or "vi " in command_lower:
        return "vim"
    elif "ssh" in command_lower:
        return "ssh"
    else:
        return "generic"
        return "ssh"


def detect_process_type(command):
    """Detect process type from command."""
    command_str = " ".join(command) if isinstance(command, list) else command
    command_lower = command_str.lower()

    if "apt" in command_lower or "apt-get" in command_lower:
        return "apt"
    elif "vim" in command_lower or "vi " in command_lower:
        return "vim"
    elif "ssh" in command_lower:
        return "ssh"
    else:
        return "generic"