"""
SSHPilot Agent - Host-side PTY handler for Flatpak sandbox

This agent runs on the host system (outside the Flatpak sandbox) and handles
proper PTY creation and shell spawning to avoid job control issues.

Inspired by the Ptyxis agent architecture.
"""

import os
import sys
import pty
import tty
import termios
import fcntl
import subprocess
import struct
import select
import signal
import pwd
import json
import logging
from typing import Optional, Tuple


def _agent_log_file() -> Optional[str]:
    """Resolve the agent's log file path: ``$XDG_STATE_HOME/sshpilot/agent.log``.

    Mirrors sshpilot's own state dir (see platform_utils.get_state_dir) so the
    agent log sits next to ``sshpilot.log`` and is easy to find. Resolved with
    the stdlib only because the agent runs standalone on the host (outside the
    Flatpak sandbox) and cannot import the sshpilot package.
    """
    try:
        state_home = os.environ.get('XDG_STATE_HOME') or os.path.join(
            os.path.expanduser('~'), '.local', 'state'
        )
        log_dir = os.path.join(state_home, 'sshpilot')
        os.makedirs(log_dir, exist_ok=True)
        return os.path.join(log_dir, 'agent.log')
    except Exception:
        return None


# Set up logging.
#
# The agent's stderr is the interactive VTE terminal in the local-terminal path,
# so logging there pollutes the shell output. Worse, once io_loop() puts the tty
# in raw mode (OPOST cleared) a log line's trailing "\n" is no longer translated
# to "\r\n", leaving the cursor mid-line and staircasing the shell prompt. Log to
# a file instead so logs never touch the terminal while staying available for
# debugging; never fall back to stderr (it may be the terminal).
_log_level = logging.DEBUG if '--verbose' in sys.argv else logging.WARNING
_log_file = _agent_log_file()
if _log_file:
    logging.basicConfig(
        level=_log_level,
        format='[sshpilot-agent] %(asctime)s %(levelname)s: %(message)s',
        filename=_log_file,
        filemode='a',
    )
else:
    logging.basicConfig(level=_log_level, handlers=[logging.NullHandler()])
logger = logging.getLogger(__name__)


class PTYAgent:
    """Agent that manages PTY creation and shell spawning on the host"""

    def __init__(self):
        self.master_fd: Optional[int] = None
        self.slave_fd: Optional[int] = None
        self.shell_pid: Optional[int] = None
        self.original_termios = None
        
    def discover_shell(self) -> str:
        """Discover the user's preferred shell on the host system"""
        try:
            # Try SHELL environment variable first
            shell = os.environ.get('SHELL')
            if shell and os.path.isfile(shell):
                logger.debug(f"Found shell from SHELL env: {shell}")
                return shell
        except Exception as e:
            logger.debug(f"Failed to get SHELL env: {e}")
        
        try:
            # Try getent passwd for current user
            username = os.environ.get('USER') or pwd.getpwuid(os.getuid()).pw_name
            result = subprocess.run(
                ['getent', 'passwd', username],
                capture_output=True,
                text=True,
                check=True
            )
            if result.stdout:
                # Format: username:x:uid:gid:gecos:home:shell
                fields = result.stdout.strip().split(':')
                if len(fields) >= 7:
                    shell = fields[6]
                    if shell and os.path.isfile(shell):
                        logger.debug(f"Found shell from getent: {shell}")
                        return shell
        except Exception as e:
            logger.debug(f"Failed to get shell from getent: {e}")
        
        try:
            # Try pwd module as fallback
            shell = pwd.getpwuid(os.getuid()).pw_shell
            if shell and os.path.isfile(shell):
                logger.debug(f"Found shell from pwd module: {shell}")
                return shell
        except Exception as e:
            logger.debug(f"Failed to get shell from pwd: {e}")
        
        # Final fallback
        logger.warning("Could not determine user shell, falling back to /bin/bash")
        return '/bin/bash'
    
    def create_pty(self) -> Tuple[int, int]:
        """
        Create a PTY master/slave pair with proper flags.
        
        This is the critical fix: we create the PTY with O_NOCTTY to prevent
        automatic controlling TTY allocation, then the shell can properly claim
        it with TIOCSCTTY.
        """
        try:
            # Create PTY pair - returns (master_fd, slave_fd)
            master_fd, slave_fd = pty.openpty()
            
            # Set O_NOCTTY on slave to prevent automatic controlling TTY
            # This allows the shell to properly claim the TTY later
            flags = fcntl.fcntl(slave_fd, fcntl.F_GETFL)
            fcntl.fcntl(slave_fd, fcntl.F_SETFL, flags | os.O_NOCTTY)
            
            logger.debug(f"Created PTY: master_fd={master_fd}, slave_fd={slave_fd}")
            
            self.master_fd = master_fd
            self.slave_fd = slave_fd
            
            return master_fd, slave_fd
            
        except Exception as e:
            logger.error(f"Failed to create PTY: {e}")
            raise
    
    def set_pty_size(self, rows: int, cols: int):
        """Set the PTY size"""
        if self.master_fd is None:
            return
        
        try:
            winsize = struct.pack('HHHH', rows, cols, 0, 0)
            fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, winsize)
            logger.debug(f"Set PTY size to {rows}x{cols}")
        except Exception as e:
            logger.error(f"Failed to set PTY size: {e}")
    
    def spawn_shell(self, shell: str, cwd: Optional[str] = None) -> int:
        """
        Spawn the user's shell with proper PTY setup.
        
        This runs the shell as a child process with the PTY slave as its
        controlling terminal, ensuring proper job control.
        """
        if self.slave_fd is None:
            raise RuntimeError("PTY not created yet")

        try:
            # Prepare environment
            env = os.environ.copy()
            # Set TERM to a proper value only if missing or set to "dumb"
            if 'TERM' not in env or env.get('TERM', '').lower() == 'dumb':
                env['TERM'] = 'xterm-256color'
            env['SHELL'] = shell
            
            # Ensure essential environment variables are set from passwd database
            # This ensures shells like zsh can properly load user configuration
            try:
                pw_entry = pwd.getpwuid(os.getuid())
                if 'USER' not in env or not env.get('USER'):
                    env['USER'] = pw_entry.pw_name
                if 'LOGNAME' not in env or not env.get('LOGNAME'):
                    env['LOGNAME'] = pw_entry.pw_name
                if 'HOME' not in env or not env.get('HOME'):
                    env['HOME'] = pw_entry.pw_dir
            except (KeyError, AttributeError):
                # If passwd lookup fails, ensure at least USER is set
                if 'USER' not in env or not env.get('USER'):
                    env['USER'] = os.getenv('USER', 'user')
                if 'LOGNAME' not in env or not env.get('LOGNAME'):
                    env['LOGNAME'] = env.get('USER', 'user')
                if 'HOME' not in env or not env.get('HOME'):
                    env['HOME'] = os.path.expanduser('~')
            
            # Set working directory
            if cwd is None:
                cwd = os.path.expanduser('~')
            
            # Fork to create shell process
            pid = os.fork()
            
            if pid == 0:
                # Child process
                try:
                    # Create new session and make this process the session leader
                    os.setsid()
                    
                    # Close master FD in child
                    os.close(self.master_fd)
                    
                    # Make the PTY slave the controlling terminal
                    # This is the critical operation that fixes job control
                    fcntl.ioctl(self.slave_fd, termios.TIOCSCTTY, 0)
                    
                    # Redirect stdin, stdout, stderr to PTY slave
                    os.dup2(self.slave_fd, 0)
                    os.dup2(self.slave_fd, 1)
                    os.dup2(self.slave_fd, 2)
                    
                    # Close the original slave FD if it's not one of stdin/stdout/stderr
                    if self.slave_fd > 2:
                        os.close(self.slave_fd)
                    
                    # Change to working directory
                    os.chdir(cwd)
                    
                    # Use interactive shell for all shells to match gnome-terminal and konsole behavior
                    # Interactive shells load user's interactive config directly (.bashrc, .zshrc, etc.)
                    # This is faster and matches what users expect from terminal emulators
                    shell_flags = ['-i']  # Interactive shell (loads interactive config files)
                    
                    # Execute the shell
                    os.execvpe(shell, [shell] + shell_flags, env)
                    
                except Exception as e:
                    logger.error(f"Child process failed: {e}")
                    sys.exit(1)
            
            else:
                # Parent process
                # Close slave FD in parent (child has its own copy)
                os.close(self.slave_fd)
                self.slave_fd = None
                
                self.shell_pid = pid
                logger.info(f"Spawned shell: {shell} (PID: {pid})")

                return pid

        except Exception as e:
            logger.error(f"Failed to spawn shell: {e}")
            raise

    def _send_status(self, status_type: str, **payload):
        """Send a structured status message to stderr for the caller."""
        stream = getattr(sys, 'stderr', None)
        if not stream:
            return

        # When the agent is spawned through VTE the stderr stream is the
        # interactive terminal itself, so emitting JSON control messages would be
        # visible to the user.  Only send structured messages when stderr is a
        # pipe (or otherwise not a tty), which covers the non-VTE execution path
        # used by the Flatpak launcher while keeping the terminal output clean.
        if hasattr(stream, 'isatty') and stream.isatty():
            return

        try:
            message = {'type': status_type}
            message.update(payload)
            stream.write(json.dumps(message) + '\n')
            stream.flush()

        except Exception as exc:
            logger.debug(f"Failed to send status message: {exc}")

    def io_loop(self):
        """
        Main I/O loop: relay data between master PTY and stdin/stdout.

        This forwards:
        - stdin (from VTE) -> master_fd (to shell)
        - master_fd (from shell) -> stdout (to VTE)
        """
        if self.master_fd is None:
            raise RuntimeError("PTY not created yet")
        
        # Set stdin to raw mode to pass through all data
        try:
            self.original_termios = termios.tcgetattr(sys.stdin.fileno())
            tty.setraw(sys.stdin.fileno())
        except Exception as e:
            logger.debug(f"Could not set raw mode on stdin: {e}")
        
        logger.debug("Starting I/O loop")
        
        try:
            while True:
                # Use select to wait for data on either stdin or master_fd
                readable, _, _ = select.select([sys.stdin, self.master_fd], [], [])
                
                for fd in readable:
                    if fd == sys.stdin:
                        # Data from VTE -> forward to shell
                        try:
                            data = os.read(sys.stdin.fileno(), 4096)
                            if not data:
                                logger.debug("stdin closed")
                                return
                            os.write(self.master_fd, data)
                        except OSError as e:
                            logger.debug(f"stdin read error: {e}")
                            return
                    
                    elif fd == self.master_fd:
                        # Data from shell -> forward to VTE
                        try:
                            data = os.read(self.master_fd, 4096)
                            if not data:
                                logger.debug("master_fd closed")
                                return
                            os.write(sys.stdout.fileno(), data)
                            sys.stdout.flush()
                        except OSError as e:
                            logger.debug(f"master_fd read error: {e}")
                            return
        
        except KeyboardInterrupt:
            logger.debug("Received interrupt")
        
        finally:
            # Restore terminal settings
            if self.original_termios:
                try:
                    termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN, self.original_termios)
                except Exception:
                    pass
    
    def cleanup(self):
        """Clean up resources"""
        logger.debug("Cleaning up agent resources")
        
        # Kill shell process if still running
        if self.shell_pid:
            try:
                os.kill(self.shell_pid, signal.SIGTERM)
                os.waitpid(self.shell_pid, 0)
            except Exception as e:
                logger.debug(f"Failed to kill shell: {e}")
        
        # Close FDs
        if self.master_fd:
            try:
                os.close(self.master_fd)
            except Exception:
                pass
        
        if self.slave_fd:
            try:
                os.close(self.slave_fd)
            except Exception:
                pass
    
    def run(self, rows: int = 24, cols: int = 80, cwd: Optional[str] = None):
        """Main entry point: create PTY, spawn shell, run I/O loop"""
        try:
            # Discover shell
            shell = self.discover_shell()
            logger.info(f"Using shell: {shell}")
            
            # Create PTY with proper flags
            self.create_pty()
            
            # Set initial size
            self.set_pty_size(rows, cols)
            
            # Spawn shell
            shell_pid = self.spawn_shell(shell, cwd)

            if shell_pid:
                self._send_status('ready', pid=shell_pid)

            # Close stderr to prevent log messages from appearing in terminal
            # unless in verbose/debug mode
            if logger.getEffectiveLevel() > logging.DEBUG:
                try:
                    # Redirect stderr to /dev/null
                    devnull = os.open(os.devnull, os.O_WRONLY)
                    os.dup2(devnull, sys.stderr.fileno())
                    os.close(devnull)
                except Exception:
                    pass
            
            # Run I/O loop (blocks until shell exits or connection closes)
            self.io_loop()
            
        except Exception as e:
            logger.error(f"Agent error: {e}", exc_info=True)
            # Send error to parent
            error_msg = json.dumps({'type': 'error', 'message': str(e)})
            sys.stderr.write(error_msg + '\n')
            sys.stderr.flush()
            return 1
        
        finally:
            self.cleanup()
        
        return 0


def handle_resize_signal(signum, frame):
    """Handle SIGWINCH for terminal resize"""
    # In a real implementation, we'd need to communicate resize events
    # from the UI to the agent, possibly via a control channel


def main():
    """Main entry point for the agent"""
    import argparse
    
    parser = argparse.ArgumentParser(description='SSHPilot PTY Agent')
    parser.add_argument('--rows', type=int, default=24, help='Terminal rows')
    parser.add_argument('--cols', type=int, default=80, help='Terminal columns')
    parser.add_argument('--cwd', type=str, default=None, help='Working directory')
    parser.add_argument('--verbose', action='store_true', help='Verbose logging')
    
    args = parser.parse_args()
    
    # Set up signal handlers
    signal.signal(signal.SIGWINCH, handle_resize_signal)
    
    # Create and run agent
    agent = PTYAgent()
    return agent.run(rows=args.rows, cols=args.cols, cwd=args.cwd)


if __name__ == '__main__':
    sys.exit(main())

