import sys
import os
import re # We need the regex module for easier parsing

from PyQt6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QHBoxLayout,
    QLineEdit, QLabel, QPushButton, QTextEdit, QMessageBox
)
from PyQt6.QtCore import Qt

class SSHConfigManager(QWidget):
    # --- Configuration for the Jump Host ---
    JUMP_HOST_NAME = "asf-jump"
    JUMP_HOST_DETAILS = f"""
Host {JUMP_HOST_NAME}
    Hostname asf-server.duckdns.org
    Port 49152
    User asf
"""

    def __init__(self):
        super().__init__()
        self.setWindowTitle("SSH Config Manager (PyQt6)")
        self.setMinimumWidth(500)
        self.config_path = os.path.expanduser("~/.ssh/config")
        self.init_ui()
        
        # Ensure directory and file exist
        os.makedirs(os.path.dirname(self.config_path), exist_ok=True)
        if not os.path.exists(self.config_path):
             with open(self.config_path, 'w') as f:
                 f.write("")

        self.ensure_jump_host_defined() # Check and add jump host immediately
        self.load_config()

    def init_ui(self):
        main_layout = QVBoxLayout()
        
        # --- 1. Input Fields for New VM ---
        input_group = QVBoxLayout()
        input_group.addWidget(QLabel("## 🛠️ Add New VM Entry"))
        
        self.host_input = QLineEdit()
        self.host_input.setPlaceholderText("Host (e.g., vm-test1)")
        
        self.port_input = QLineEdit()
        self.port_input.setPlaceholderText("Port (e.g., 6002)")
        
        self.user_input = QLineEdit()
        self.user_input.setPlaceholderText("User (e.g., asf_user)")
        
        self.proxy_input = QLineEdit()
        self.proxy_input.setPlaceholderText(f"ProxyJump (Default: {self.JUMP_HOST_NAME})")
        self.proxy_input.setText(self.JUMP_HOST_NAME) # Pre-fill the jump host name
        
        input_group.addWidget(self.host_input)
        input_group.addWidget(self.port_input)
        input_group.addWidget(self.user_input)
        input_group.addWidget(self.proxy_input)
        
        add_button = QPushButton("➕ Add VM & Save Config")
        add_button.clicked.connect(self.add_vm_entry)
        input_group.addWidget(add_button)
        
        main_layout.addLayout(input_group)
        main_layout.addWidget(QLabel("---"))

        # --- 2. Existing Config Viewer ---
        main_layout.addWidget(QLabel("## 📜 Current ~/.ssh/config Content"))
        self.config_viewer = QTextEdit()
        self.config_viewer.setReadOnly(True)
        self.config_viewer.setMinimumHeight(150)
        main_layout.addWidget(self.config_viewer)
        
        # --- 3. Connection Instructions & Display ---
        main_layout.addWidget(QLabel("## 🔗 Connection Info"))
        
        info_layout = QHBoxLayout()
        info_layout.addWidget(QLabel("Select VM Host:"))
        self.vm_select = QLineEdit()
        self.vm_select.setPlaceholderText("Enter Host name (e.g., vm-test1)")
        info_layout.addWidget(self.vm_select)
        
        show_button = QPushButton("Show Commands")
        show_button.clicked.connect(self.show_connection_commands)
        info_layout.addWidget(show_button)
        
        main_layout.addLayout(info_layout)

        self.command_output = QTextEdit()
        self.command_output.setReadOnly(True)
        self.command_output.setMinimumHeight(100)
        main_layout.addWidget(self.command_output)

        self.setLayout(main_layout)

    def load_config(self):
        """Loads the current config file content into the viewer."""
        try:
            with open(self.config_path, 'r') as f:
                content = f.read()
                self.config_viewer.setText(content)
        except Exception as e:
            self.config_viewer.setText(f"Error loading config: {e}")

    def ensure_jump_host_defined(self):
        """Checks if the JUMP_HOST_NAME is defined and adds it if missing."""
        try:
            with open(self.config_path, 'r+') as f:
                content = f.read()
                
                # Use regex to find if the Host definition already exists
                if re.search(rf"^Host\s+{re.escape(self.JUMP_HOST_NAME)}\s*$", content, re.MULTILINE | re.IGNORECASE):
                    # print("Jump host already defined.") # For debugging
                    return

                # If not found, append the definition to the beginning of the file
                f.seek(0)
                f.write(self.JUMP_HOST_DETAILS + "\n" + content)
                f.truncate()
                
                QMessageBox.information(self, "Setup Complete", 
                                        f"Automatically added the '{self.JUMP_HOST_NAME}' jump host definition to the config file.")

        except Exception as e:
            QMessageBox.critical(self, "Setup Error", f"Failed to ensure jump host definition: {e}")

    def add_vm_entry(self):
        """Validates inputs, formats the entry, and appends it to the config file."""
        host = self.host_input.text().strip()
        port = self.port_input.text().strip()
        user = self.user_input.text().strip()
        proxy = self.proxy_input.text().strip()

        if not all([host, port, user, proxy]):
            QMessageBox.warning(self, "Input Error", "All fields must be filled out.")
            return
        
        # Ensure jump host is defined before adding dependent VMs
        self.ensure_jump_host_defined() 

        new_entry = f"""

Host {host}
    Hostname 127.0.0.1
    Port {port}
    User {user}
    ProxyJump {proxy}
"""

        try:
            with open(self.config_path, 'a') as f:
                f.write(new_entry)
            
            QMessageBox.information(self, "Success", f"VM '{host}' successfully added to {self.config_path}!")
            self.load_config() # Refresh the viewer
            
            # Clear inputs
            self.host_input.clear()
            self.port_input.clear()
            self.user_input.clear()
            self.proxy_input.setText(self.JUMP_HOST_NAME) # Reset ProxyJump field

        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to save file: {e}")

    def show_connection_commands(self):
        """Generates and displays the connection commands for the selected VM."""
        host = self.vm_select.text().strip()
        
        if not host:
            self.command_output.setText("Please enter a Host name to look up.")
            return

        ssh_port = None
        try:
            with open(self.config_path, 'r') as f:
                content = f.read().splitlines()
            
            # Simple line-by-line parser to find the port
            current_host = None
            for line in content:
                line = line.strip()
                if line.startswith('Host '):
                    current_host = line.split()[1]
                elif current_host == host and line.startswith('Port '):
                    ssh_port = line.split()[1]
                    break
        except Exception:
            self.command_output.setText(f"Error reading config file to find port for {host}.")
            return
            
        if not ssh_port:
            self.command_output.setText(f"Host '{host}' not found or Port not specified in config.")
            return

        # VNC Port (Assuming VNC Port = SSH Port - 1000, based on your ranges 6002/5002)
        vnc_port = int(ssh_port) - 1000 
        
        # 2. Generate Commands
        commands = f"""
✅ **SSH Command** (for direct console access via Jump Host):
ssh {host}

🖥️ **VNC Command** (for secure graphical access via SSH Tunnel):
1. Establish the tunnel (keep this window open/running):
   ssh -L 5900:127.0.0.1:{vnc_port} {host} -N &

2. Connect VNC Client to your local machine:
   VNC Host/Port: 127.0.0.1:5900
"""
        self.command_output.setText(commands)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    manager = SSHConfigManager()
    manager.show()
    sys.exit(app.exec())