rp/pr/tools/process_handlers.py

import re
import time
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'