import os
import re
import subprocess
import sys

def get_ssh_history():
    """Reads ~/.bash_history and extracts unique SSH host strings."""
    # NOTE: This assumes Bash history. Adjust for Zsh (~/.zsh_history) if needed.
    history_file = os.path.expanduser("~/.bash_history")
    
    # Regex to capture the host part of an SSH command: 'ssh user@host [options]'
    SSH_REGEX = re.compile(r"^\s*ssh\s+([^-@\s]+(?:@\S+)?)") 
    
    hosts = []
    
    try:
        # Read history file in reverse (newest first)
        with open(history_file, 'r') as f:
            for line in reversed(f.readlines()):
                match = SSH_REGEX.search(line)
                if match:
                    host_string = match.group(1).split()[0] # Take only the first word after 'ssh'
                    
                    # Deduplicate the list
                    if host_string not in hosts:
                        hosts.append(host_string)
    except FileNotFoundError:
        print(f"Error: History file not found at {history_file}", file=sys.stderr)
        sys.exit(1)
        
    return hosts

def run_fzf_selection(hosts):
    """Pipes the list of hosts to fzf and returns the selected host."""
    
    if not hosts:
        print("No SSH commands found in history.", file=sys.stderr)
        return None

    try:
        # 1. Join hosts with newlines for piping
        input_data = "\n".join(hosts)
        
        # 2. Call fzf via subprocess
        # --no-sort ensures the history order (newest first) is maintained
        # --prompt sets a helpful prompt text
        process = subprocess.run(
            ['fzf', '--no-sort', '--prompt=Select SSH Host: '],
            input=input_data,
            capture_output=True,
            text=True,
            check=True # Raise an exception for non-zero exit codes (e.g., fzf not installed)
        )
        
        # fzf prints the selected line to stdout
        return process.stdout.strip()
        
    except FileNotFoundError:
        print("Error: 'fzf' is not installed or not in your PATH.", file=sys.stderr)
        print("Please install fzf (e.g., 'brew install fzf' or 'sudo apt install fzf').", file=sys.stderr)
        sys.exit(1)
    except subprocess.CalledProcessError:
        # User canceled fzf (e.g., pressed Esc)
        sys.exit(0)

def main():
    history_hosts = get_ssh_history()
    selected_host = run_fzf_selection(history_hosts)
    
    if selected_host:
        print(f"Connecting to: {selected_host}...")
        
        # Crucial step: Use os.execvp to replace the current Python process 
        # with the SSH process, passing control to the SSH client.
        try:
            os.execvp("ssh", ["ssh", selected_host])
        except FileNotFoundError:
            print("Error: 'ssh' command not found.", file=sys.stderr)
            sys.exit(1)
        
if __name__ == "__main__":
    main()