"""
SSH Tab for Infra TUI
"""

import os
import subprocess
import tempfile
from textual.app import ComposeResult
from textual.widgets import Static, Label, Button, DataTable
from textual.containers import Horizontal
from ..core.users import UserManager
from .dialogs import ConfirmDialog, SSHTestDialog, SSHKeySetupDialog


class SSHTab(Static):
    """Tab for SSH security management"""

    # State tracking
    selected_user: str | None = None
    test_results: dict = {}  # {user: (success, method)}
    user_data: dict = {}  # {user: {has_key, has_password, is_sudo}}

    def compose(self) -> ComposeResult:
        yield Label("SSH Security Management", id="ssh-title")
        yield Label("Press **R** to refresh", id="ssh-help")

        # Global Settings Section
        yield Label("Global Settings", id="ssh-config-title")
        yield DataTable(id="ssh-config-table")
        yield Horizontal(
            Button("Disable SSH Password", id="ssh-disable-password", variant="error"),
            Button("Disable Root Login", id="ssh-disable-root", variant="error"),
            id="ssh-global-buttons"
        )

        yield Label("", id="ssh-spacer")

        # User SSH Access Section
        yield Label("User SSH Access (select row to manage)", id="ssh-users-title")
        yield DataTable(id="ssh-users-table", cursor_type="row")
        yield Label("Selected: [dim]None[/]", id="ssh-selected-label")
        yield Horizontal(
            Button("Test SSH", id="ssh-test", variant="primary", disabled=True),
            Button("Disable SSH Pass", id="ssh-lock-password", variant="warning", disabled=True),
            Button("Disable SSH", id="ssh-disable-user", variant="error", disabled=True),
            Button("Setup Key", id="ssh-setup-keys", variant="success", disabled=True),
            id="ssh-user-buttons"
        )

    def on_mount(self) -> None:
        """Initialize SSH tables when tab is mounted"""
        self.refresh_ssh_config()
        self.refresh_ssh_users()

    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        """Handle row selection in the users table"""
        if event.data_table.id != "ssh-users-table":
            return

        # Get the username from the row key
        self.selected_user = str(event.row_key.value) if event.row_key else None
        self.update_selected_label()
        self.update_button_states()

    def update_selected_label(self) -> None:
        """Update the selected user label"""
        try:
            label = self.query_one("#ssh-selected-label", Label)
            if self.selected_user:
                user_info = self.user_data.get(self.selected_user, {})
                sudo_text = " [cyan](sudo)[/]" if user_info.get("is_sudo") else ""
                test_result = self.test_results.get(self.selected_user)
                if test_result:
                    success, method = test_result
                    test_text = f" - [green]Working ({method})[/]" if success else f" - [red]Failed[/]"
                else:
                    test_text = " - [dim]Not tested[/]"
                label.update(f"Selected: [bold]{self.selected_user}[/]{sudo_text}{test_text}")
            else:
                label.update("Selected: [dim]None[/]")
        except Exception:
            pass

    def update_button_states(self) -> None:
        """Enable/disable buttons based on selected user state"""
        try:
            test_btn = self.query_one("#ssh-test", Button)
            lock_btn = self.query_one("#ssh-lock-password", Button)
            disable_btn = self.query_one("#ssh-disable-user", Button)
            setup_btn = self.query_one("#ssh-setup-keys", Button)

            if not self.selected_user:
                test_btn.disabled = True
                lock_btn.disabled = True
                disable_btn.disabled = True
                setup_btn.disabled = True
                return

            # Enable all buttons when user is selected
            test_btn.disabled = False
            disable_btn.disabled = False
            setup_btn.disabled = False

            # Check if user already has SSH password disabled via Match block
            ssh_pass_disabled = self._user_has_ssh_password_disabled(self.selected_user)

            if ssh_pass_disabled:
                lock_btn.disabled = True
                lock_btn.label = "SSH Pass Disabled"
            else:
                lock_btn.disabled = False
                lock_btn.label = "Disable SSH Pass"

        except Exception:
            pass

    def _user_has_ssh_password_disabled(self, user: str) -> bool:
        """Check if user has SSH password disabled via Match block in sshd_config"""
        try:
            config_path = "/etc/ssh/sshd_config"
            if not os.path.exists(config_path):
                return False

            with open(config_path) as f:
                in_match_block = False
                for line in f:
                    stripped = line.strip()
                    if stripped.startswith("Match User ") and user in stripped:
                        in_match_block = True
                        continue
                    if in_match_block:
                        if "PasswordAuthentication no" in stripped:
                            return True
                        if not stripped or stripped.startswith("#") or stripped.startswith("Match "):
                            break
            return False
        except Exception:
            return False

    def refresh_ssh_config(self) -> None:
        """Refresh SSH configuration table"""
        try:
            import subprocess
            table = self.query_one("#ssh-config-table", DataTable)
            table.clear(columns=True)
            table.add_column("Setting", width=30)
            table.add_column("Value", width=20)
            table.add_column("Status", width=30)

            # Read sshd_config
            config_path = "/etc/ssh/sshd_config"
            settings_to_check = {
                "PermitRootLogin": ["no", "without-password", "prohibit-password"],
                "PasswordAuthentication": ["no"],
                "PubkeyAuthentication": ["yes"],
                "PermitEmptyPasswords": ["no"],
            }

            config = {}
            if os.path.exists(config_path):
                with open(config_path) as f:
                    for line in f:
                        line = line.strip()
                        if line and not line.startswith("#") and " " in line:
                            key, value = line.split(None, 1)
                            config[key] = value.strip()  # Strip whitespace from value

            for setting, recommended in settings_to_check.items():
                current = config.get(setting, "not set")
                if current in recommended:
                    status = "[green]Secure[/]"
                elif current == "not set":
                    status = "[yellow]Not Set[/]"
                elif setting == "PermitRootLogin" and current == "yes":
                    status = "[red]Root can login[/]"
                elif setting == "PasswordAuthentication" and current == "yes":
                    status = "[red]Passwords allowed[/]"
                else:
                    status = "[yellow]Review[/]"

                table.add_row(setting, current, status)

            # Update global button states
            self.update_global_button_states(config)
        except Exception:
            pass

    def update_global_button_states(self, config: dict) -> None:
        """Update global action button states based on config"""
        try:
            disable_pass_btn = self.query_one("#ssh-disable-password", Button)
            disable_root_btn = self.query_one("#ssh-disable-root", Button)

            # Disable password auth button if already disabled (SSH global setting)
            pass_auth = config.get("PasswordAuthentication", "yes").strip()
            if pass_auth == "no":
                disable_pass_btn.disabled = True
                disable_pass_btn.label = "SSH Password Disabled"
            else:
                disable_pass_btn.disabled = False
                disable_pass_btn.label = "Disable SSH Password"

            # Disable root login button if already disabled
            root_login = config.get("PermitRootLogin", "yes").strip()
            if root_login in ["no", "without-password", "prohibit-password"]:
                disable_root_btn.disabled = True
                disable_root_btn.label = "Root Login Disabled"
            else:
                disable_root_btn.disabled = False
                disable_root_btn.label = "Disable Root Login"
        except Exception:
            pass

    def refresh_ssh_users(self) -> None:
        """Refresh SSH user access table"""
        try:
            import subprocess
            table = self.query_one("#ssh-users-table", DataTable)
            table.clear(columns=True)
            table.add_column("User", width=15, key="user")
            table.add_column("SSH Key", width=12)
            table.add_column("Password", width=12)
            table.add_column("Last Test", width=25)

            all_users = UserManager.get_all_users()

            # System users to exclude (but keep root)
            system_users = {
                "nobody", "daemon", "bin", "sys", "sync", "games", "man",
                "lp", "mail", "news", "uucp", "proxy", "www-data", "backup", "list",
                "irc", "gnats", "_apt", "systemd-timesync", "systemd-network",
                "systemd-resolve", "messagebus", "sshd"
            }

            # Include root explicitly
            users_to_show = ["root"] + sorted([u for u in all_users if u not in system_users and u != "root"])

            self.user_data = {}

            for user in users_to_show:
                # Check for SSH keys
                if user == "root":
                    ssh_dir = "/root/.ssh"
                else:
                    ssh_dir = f"/home/{user}/.ssh"

                has_key = False
                if os.path.exists(ssh_dir):
                    authorized_keys = os.path.join(ssh_dir, "authorized_keys")
                    if os.path.exists(authorized_keys):
                        try:
                            with open(authorized_keys) as f:
                                if f.read().strip():
                                    has_key = True
                        except PermissionError:
                            pass

                # Check password status (account-wide via passwd -S)
                try:
                    result = subprocess.run(
                        ["sudo", "-n", "passwd", "-S", user],
                        capture_output=True,
                        text=True
                    )
                    parts = result.stdout.split()
                    account_password_enabled = len(parts) >= 2 and parts[1] == "P"
                except Exception:
                    account_password_enabled = True  # Assume enabled if check fails

                # Check if SSH password is disabled via Match User block
                ssh_password_disabled = self._user_has_ssh_password_disabled(user)

                # Password status: show as disabled if SSH password is disabled
                has_password = account_password_enabled and not ssh_password_disabled

                # Check if user is sudoer
                is_sudo = user == "root" or UserManager.is_sudoer(user)

                # Store user data for button state logic
                self.user_data[user] = {
                    "has_key": has_key,
                    "has_password": has_password,
                    "is_sudo": is_sudo
                }

                # Format display
                key_status = "[green]Yes[/]" if has_key else "[red]No[/]"

                # Password status: show different states
                if ssh_password_disabled:
                    pass_status = "[yellow]SSH Disabled[/]"
                elif not account_password_enabled:
                    pass_status = "[red]Locked[/]"
                else:
                    pass_status = "[green]Enabled[/]"

                # Get test result if available
                test_result = self.test_results.get(user)
                if test_result:
                    success, method = test_result
                    test_status = f"[green]{method}[/]" if success else "[red]Failed[/]"
                else:
                    test_status = "[dim]Not tested[/]"

                table.add_row(user, key_status, pass_status, test_status, key=user)

            # Reset selection state
            self.selected_user = None
            self.update_selected_label()
            self.update_button_states()

        except Exception:
            pass

    # === SSH Action Methods ===

    def disable_password_auth(self) -> None:
        """Disable password authentication for SSH"""
        # Count users with SSH keys
        users_with_keys = 0
        all_users = UserManager.get_all_users()

        system_users = {
            "root", "nobody", "daemon", "bin", "sys", "sync", "games", "man",
            "lp", "mail", "news", "uucp", "proxy", "www-data", "backup", "list",
            "irc", "gnats", "_apt", "systemd-timesync", "systemd-network",
            "systemd-resolve", "messagebus", "sshd"
        }

        for user in all_users:
            if user in system_users:
                continue
            ssh_dir = f"/home/{user}/.ssh"
            if os.path.exists(ssh_dir):
                authorized_keys = os.path.join(ssh_dir, "authorized_keys")
                if os.path.exists(authorized_keys):
                    with open(authorized_keys) as f:
                        if f.read().strip():
                            users_with_keys += 1

        if users_with_keys < 1:
            self.app.notify(
                f"Cannot disable password auth: Only {users_with_keys} user(s) with SSH keys. "
                "Setup SSH keys first for at least one user.",
                severity="error"
            )
            return

        def do_disable():
            # Update sshd_config
            config_path = "/etc/ssh/sshd_config"
            config = []

            with open(config_path) as f:
                for line in f:
                    stripped = line.strip()
                    if stripped.startswith("PasswordAuthentication"):
                        # Comment out the existing line
                        config.append(f"# {line}")
                        config.append("PasswordAuthentication no\n")
                    elif stripped.startswith("#PasswordAuthentication no"):
                        # Already disabled, skip
                        config.append(line)
                    else:
                        config.append(line)

            # Write backup
            subprocess.run(["sudo", "cp", config_path, config_path + ".bak"], check=False)

            # Write new config
            with open("/tmp/sshd_config", "w") as f:
                f.writelines(config)

            subprocess.run(["sudo", "mv", "/tmp/sshd_config", config_path], check=True)
            subprocess.run(["sudo", "systemctl", "reload", "sshd"], check=True)

            self.app.notify("Password authentication disabled. SSH reloaded.", severity="information")
            self.refresh_ssh_config()

        self.app.push_screen(ConfirmDialog(
            f"**Disable SSH password authentication?**\n\n"
            f"This will disable password-based SSH login.\n"
            f"Users will need SSH keys to access the server.\n\n"
            f"Users with SSH keys: {users_with_keys}\n\n"
            f"Make sure you have SSH keys setup before proceeding!",
            on_confirm=do_disable
        ))

    def setup_ssh_keys(self) -> None:
        """Setup SSH keys for a user - show dialog to paste public key"""
        user = self.selected_user

        if not user:
            self.app.notify("No user selected", severity="error")
            return

        def on_dialog_result(result):
            if result:
                self._add_ssh_key(result["user"], result["key"])

        self.app.push_screen(SSHKeySetupDialog(self.app, user), on_dialog_result)

    def _add_ssh_key(self, user: str, key: str) -> None:
        """Add SSH public key for a user with proper permissions"""
        if user == "root":
            ssh_dir = "/root/.ssh"
            home_dir = "/root"
        else:
            home_dir = f"/home/{user}"
            ssh_dir = f"{home_dir}/.ssh"

        auth_keys = f"{ssh_dir}/authorized_keys"

        try:
            # Create .ssh directory if needed
            if not os.path.exists(ssh_dir):
                subprocess.run(["sudo", "mkdir", "-p", ssh_dir], check=True)

            # Add key to authorized_keys (append)
            # Use a temp file and sudo to append
            with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
                f.write(key + "\n")
                temp_file = f.name

            # Append to authorized_keys
            subprocess.run(
                f"sudo cat {temp_file} | sudo tee -a {auth_keys} > /dev/null",
                shell=True, check=True
            )
            os.unlink(temp_file)

            # Set correct ownership and permissions
            subprocess.run(["sudo", "chown", "-R", f"{user}:{user}", ssh_dir], check=True)
            subprocess.run(["sudo", "chmod", "700", ssh_dir], check=True)
            subprocess.run(["sudo", "chmod", "600", auth_keys], check=True)

            self.app.notify(f"SSH key added for {user}", severity="information")

            # Refresh the SSH users table
            self.refresh_ssh_users()
            self.selected_user = user
            self.update_selected_label()
            self.update_button_states()

        except subprocess.CalledProcessError as e:
            self.app.notify(f"Failed to add SSH key: {e}", severity="error")
        except Exception as e:
            self.app.notify(f"Error: {e}", severity="error")

    def _get_ssh_host(self) -> str:
        """Get the best IP/hostname for SSH testing. Prefer external IPv4."""
        import requests

        # Try to get external IPv4 from ipinfo.io
        try:
            resp = requests.get("https://ipinfo.io/ip", timeout=5)
            if resp.status_code == 200:
                ip = resp.text.strip()
                if ip and ":" not in ip:  # IPv4 (no colons)
                    return ip
        except Exception:
            pass

        # Try curl ifconfig.me as fallback
        try:
            resp = requests.get("https://ifconfig.me", timeout=5)
            if resp.status_code == 200:
                ip = resp.text.strip()
                if ip and ":" not in ip:
                    return ip
        except Exception:
            pass

        # Fallback to hostname -I (first IP)
        try:
            result = subprocess.run(
                ["hostname", "-I"],
                capture_output=True,
                text=True,
                timeout=5
            )
            if result.returncode == 0:
                ips = result.stdout.strip().split()
                # Prefer IPv4 addresses
                for ip in ips:
                    if ":" not in ip:  # IPv4
                        return ip
                if ips:
                    return ips[0]
        except Exception:
            pass

        # Last resort
        return "localhost"

    def test_ssh_for_user(self) -> None:
        """Test SSH connection for the selected user using paramiko"""
        import paramiko

        user = self.selected_user

        if not user:
            self.app.notify("No user selected", severity="error")
            return

        self.app.notify(f"Testing SSH for {user}...", severity="information")

        def do_initial_test():
            host = self._get_ssh_host()
            self.app.call_from_thread(
                self.app.notify, f"Testing SSH to {host}...", severity="information"
            )

            # Try key-based auth without passphrase first
            key_auth_worked = False
            try:
                client = paramiko.SSHClient()
                client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

                client.connect(
                    hostname=host,
                    username=user,
                    timeout=10,
                    look_for_keys=True,
                    allow_agent=True,
                    password=None
                )
                # Test connection
                stdin, stdout, stderr = client.exec_command("echo ok")
                output = stdout.read().decode().strip()
                client.close()

                if output == "ok":
                    key_auth_worked = True
                    self.test_results[user] = (True, "key")
                    self.app.call_from_thread(
                        self.app.notify, f"SSH working for {user}@{host} (key auth)", severity="information"
                    )
                    self._update_ssh_ui_after_test(user)

            except Exception as e:
                # Key auth failed for any reason - close client and show dialog
                try:
                    client.close()
                except Exception:
                    pass

            # If key auth didn't work, show dialog for password/passphrase
            if not key_auth_worked:
                self.app.call_from_thread(
                    self._show_ssh_test_dialog, user, host, True, True
                )

        self.app.run_worker(do_initial_test, thread=True)

    def _show_ssh_test_dialog(self, user: str, host: str, has_key: bool, has_password: bool) -> None:
        """Show the SSH test dialog for credential input"""
        def on_dialog_result(result):
            if result:
                result["host"] = host
                self.app.run_worker(
                    lambda: self._test_ssh_with_credential(result),
                    thread=True
                )

        self.app.push_screen(
            SSHTestDialog(self.app, user, has_key, has_password),
            on_dialog_result
        )

    def _test_ssh_with_credential(self, cred_info: dict) -> None:
        """Test SSH with provided credential using paramiko"""
        import paramiko

        user = cred_info["user"]
        host = cred_info.get("host", "localhost")
        method = cred_info["method"]

        client = None
        temp_key_file = None

        try:
            client = paramiko.SSHClient()
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            if method == "key":
                # Use provided private key content or path
                key_content = cred_info.get("key_content", "")
                passphrase = cred_info.get("passphrase")

                # Check if key_content is a file path or actual key content
                if os.path.exists(key_content):
                    # It's a file path - read it
                    key_path = key_content
                else:
                    # It's key content - write to temp file
                    with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.pem') as f:
                        f.write(key_content)
                        key_path = f.name
                        temp_key_file = key_path

                # Try to load the key with optional passphrase
                pkey = None
                key_errors = []

                # Try different key types
                key_classes = [
                    ("Ed25519", paramiko.Ed25519Key),
                    ("RSA", paramiko.RSAKey),
                    ("ECDSA", paramiko.ECDSAKey),
                ]

                for key_type, key_class in key_classes:
                    try:
                        if passphrase:
                            pkey = key_class.from_private_key_file(key_path, password=passphrase)
                        else:
                            pkey = key_class.from_private_key_file(key_path)
                        break  # Successfully loaded
                    except paramiko.ssh_exception.PasswordRequiredException:
                        key_errors.append(f"{key_type}: requires passphrase")
                    except paramiko.ssh_exception.SSHException:
                        # Not this key type, try next
                        continue
                    except Exception as e:
                        key_errors.append(f"{key_type}: {e}")

                if not pkey:
                    error_msg = f"Failed to load key. Try: " + "; ".join(key_errors[:3])
                    self.test_results[user] = (False, "invalid key")
                    self.app.call_from_thread(
                        self.app.notify, f"{error_msg}", severity="error"
                    )
                    self._update_ssh_ui_after_test(user)
                    return

                # Connect with the key
                client.connect(
                    hostname=host,
                    username=user,
                    pkey=pkey,
                    timeout=10
                )

            elif method == "password":
                # Use password auth
                password = cred_info.get("credential", "")
                client.connect(
                    hostname=host,
                    username=user,
                    password=password,
                    timeout=10,
                    look_for_keys=False,
                    allow_agent=False
                )
            else:
                raise Exception(f"Unknown method: {method}")

            # Test connection
            stdin, stdout, stderr = client.exec_command("echo ok")
            output = stdout.read().decode().strip()
            client.close()
            client = None

            if output == "ok":
                method_text = "key+passphrase" if method == "key" and cred_info.get("passphrase") else method
                self.test_results[user] = (True, method_text)
                self.app.call_from_thread(
                    self.app.notify, f"SSH working for {user}@{host} ({method_text})", severity="information"
                )
            else:
                self.test_results[user] = (False, f"{method} failed")
                self.app.call_from_thread(
                    self.app.notify, f"SSH {method} auth failed", severity="error"
                )

        except paramiko.AuthenticationException as e:
            self.test_results[user] = (False, "auth failed")
            self.app.call_from_thread(
                self.app.notify, f"SSH authentication failed. Check key/passphrase/password.", severity="error"
            )
        except paramiko.SSHException as e:
            self.test_results[user] = (False, "ssh error")
            self.app.call_from_thread(
                self.app.notify, f"SSH error: {e}", severity="error"
            )
        except Exception as e:
            self.test_results[user] = (False, f"error: {e}")
            self.app.call_from_thread(
                self.app.notify, f"SSH test error: {e}", severity="error"
            )
        finally:
            # Clean up
            if client:
                try:
                    client.close()
                except Exception:
                    pass
            if temp_key_file and os.path.exists(temp_key_file):
                try:
                    os.unlink(temp_key_file)
                except Exception:
                    pass

        self._update_ssh_ui_after_test(user)

    def _update_ssh_ui_after_test(self, user: str) -> None:
        """Update SSH tab UI after a test completes"""
        try:
            self.app.call_from_thread(self.refresh_ssh_users)
            self.selected_user = user
            self.app.call_from_thread(self.update_selected_label)
            self.app.call_from_thread(self.update_button_states)
        except Exception:
            pass

    def lock_user_password(self) -> None:
        """Disable SSH password auth for the selected user (SSH only, not account-wide)"""
        user = self.selected_user

        if not user:
            self.app.notify("No user selected", severity="error")
            return

        # Safety check - at least one user must have SSH key access (tested or configured)
        users_with_ssh = []

        # First check users with verified/tested SSH
        for other_user, other_result in self.test_results.items():
            if other_result and other_result[0]:
                method = other_result[1]
                if method in ("key", "key+passphrase"):
                    users_with_ssh.append(f"{other_user} (tested)")

        # Also check users with SSH keys configured (even if not tested)
        for other_user, info in self.user_data.items():
            if info.get("has_key") and other_user not in [u.split(" (")[0] for u in users_with_ssh]:
                users_with_ssh.append(f"{other_user} (has key)")

        if not users_with_ssh:
            self.app.notify(
                "Cannot disable SSH password: No user has SSH key access.\n"
                "Setup SSH keys for at least one user first.",
                severity="error"
            )
            return

        def do_lock():
            try:
                config_path = "/etc/ssh/sshd_config"
                config_lines = []
                match_block_added = False

                with open(config_path) as f:
                    in_match_block = False
                    for line in f:
                        stripped = line.strip()

                        # Skip existing Match block for this user
                        if stripped.startswith("Match User ") and user in stripped:
                            in_match_block = True
                            continue

                        # End of Match block
                        if in_match_block and (not stripped or stripped.startswith("#") or stripped.startswith("Match ")):
                            if stripped.startswith("Match "):
                                # New Match block, add our block before this
                                config_lines.append(f"\n# SSH password disabled for {user}\n")
                                config_lines.append(f"Match User {user}\n")
                                config_lines.append(f"    PasswordAuthentication no\n")
                                match_block_added = True
                            in_match_block = False
                            if not stripped.startswith("Match "):
                                config_lines.append(line)
                            elif stripped.startswith("Match "):
                                config_lines.append(line)
                            continue

                        if in_match_block:
                            continue  # Skip lines inside the Match block we're removing

                        config_lines.append(line)

                # If we didn't find an existing Match block, add one at the end
                if not match_block_added:
                    config_lines.append(f"\n# SSH password disabled for {user}\n")
                    config_lines.append(f"Match User {user}\n")
                    config_lines.append(f"    PasswordAuthentication no\n")

                # Write backup and new config
                subprocess.run(["sudo", "cp", config_path, config_path + ".bak"], check=False)

                with open("/tmp/sshd_config", "w") as f:
                    f.writelines(config_lines)

                subprocess.run(["sudo", "mv", "/tmp/sshd_config", config_path], check=True)
                subprocess.run(["sudo", "systemctl", "reload", "sshd"], check=True)

                self.app.notify(f"SSH password disabled for {user} (account password still works for RDP, console)", severity="information")
                self.refresh_ssh_config()
                self.refresh_ssh_users()
                self.selected_user = user
                self.update_selected_label()
                self.update_button_states()

            except subprocess.CalledProcessError as e:
                self.app.notify(f"Failed to disable SSH password: {e}", severity="error")
            except Exception as e:
                self.app.notify(f"Error: {e}", severity="error")

        self.app.push_screen(ConfirmDialog(
            f"**Disable SSH password for {user}?**\n\n"
            f"This will disable password authentication for SSH only.\n"
            f"The account password will still work for:\n"
            f"* RDP (Remote Desktop)\n"
            f"* Console login\n"
            f"* Other services\n\n"
            f"Users with SSH key access:\n" + "\n".join(f"* {u}" for u in users_with_ssh),
            on_confirm=do_lock
        ))

    def disable_user_ssh(self) -> None:
        """Disable all SSH access for the selected user"""
        user = self.selected_user

        if not user:
            self.app.notify("No user selected", severity="error")
            return

        # Safety check - ensure another user has SSH access (tested or configured)
        other_users_with_ssh = []

        # Check users with verified/tested SSH
        for other_user, other_result in self.test_results.items():
            if other_user != user and other_result and other_result[0]:
                method = other_result[1]
                if method in ("key", "key+passphrase"):
                    other_users_with_ssh.append(f"{other_user} (tested)")

        # Also check users with SSH keys configured
        for other_user, info in self.user_data.items():
            if other_user != user and info.get("has_key"):
                user_display = f"{other_user} (has key)"
                if user_display not in other_users_with_ssh:
                    other_users_with_ssh.append(user_display)

        if not other_users_with_ssh:
            self.app.notify(
                "Cannot disable SSH: No other user has SSH key access.\n"
                "Setup SSH keys for another user first to avoid lockout.",
                severity="error"
            )
            return

        user_info = self.user_data.get(user, {})
        has_key = user_info.get("has_key", False)
        has_password = user_info.get("has_password", False)

        actions = []
        if has_key:
            actions.append("Remove SSH authorized_keys")
        if has_password:
            actions.append("Lock password")

        if not actions:
            self.app.notify(f"{user} already has no SSH access", severity="information")
            return

        def do_disable():
            try:
                # Lock password if enabled
                if has_password:
                    subprocess.run(["sudo", "passwd", "-l", user], check=True, capture_output=True)

                # Remove/rename authorized_keys if exists
                if has_key:
                    if user == "root":
                        auth_keys = "/root/.ssh/authorized_keys"
                    else:
                        auth_keys = f"/home/{user}/.ssh/authorized_keys"

                    if os.path.exists(auth_keys):
                        # Rename instead of delete (safer, can recover)
                        backup = auth_keys + ".disabled"
                        subprocess.run(["sudo", "mv", auth_keys, backup], check=True, capture_output=True)

                self.app.notify(f"SSH access disabled for {user}", severity="information")
                self.test_results[user] = (False, "disabled")
                self.refresh_ssh_users()
                self.selected_user = user
                self.update_selected_label()
                self.update_button_states()

            except subprocess.CalledProcessError as e:
                self.app.notify(f"Failed to disable SSH: {e}", severity="error")

        self.app.push_screen(ConfirmDialog(
            f"**Disable ALL SSH access for {user}?**\n\n"
            f"This will:\n"
            + "\n".join(f"* {a}" for a in actions) +
            f"\n\nOther users with SSH access:\n" + "\n".join(f"* {u}" for u in other_users_with_ssh),
            on_confirm=do_disable
        ))

    def disable_root_login(self) -> None:
        """Disable root SSH login"""
        # Check if at least one non-root sudo user has SSH access (tested or configured)
        sudo_users_with_ssh = []

        for user, info in self.user_data.items():
            if user == "root":
                continue
            if info.get("is_sudo") and info.get("has_key"):
                # Check if also tested
                test_result = self.test_results.get(user)
                if test_result and test_result[0]:
                    sudo_users_with_ssh.append(f"{user} (tested)")
                else:
                    sudo_users_with_ssh.append(f"{user} (has key)")

        if not sudo_users_with_ssh:
            self.app.notify(
                "Cannot disable root login: No non-root sudo user has SSH key access.\n"
                "Setup SSH keys for a sudo user first.",
                severity="error"
            )
            return

        def do_disable():
            config_path = "/etc/ssh/sshd_config"
            config_lines = []

            try:
                with open(config_path) as f:
                    found = False
                    for line in f:
                        stripped = line.strip()
                        if stripped.startswith("PermitRootLogin"):
                            config_lines.append(f"# {line}")
                            config_lines.append("PermitRootLogin no\n")
                            found = True
                        else:
                            config_lines.append(line)

                    if not found:
                        config_lines.append("\n# Disable root login\nPermitRootLogin no\n")

                # Write backup and new config
                subprocess.run(["sudo", "cp", config_path, config_path + ".bak"], check=False)

                with open("/tmp/sshd_config", "w") as f:
                    f.writelines(config_lines)

                subprocess.run(["sudo", "mv", "/tmp/sshd_config", config_path], check=True)
                subprocess.run(["sudo", "systemctl", "reload", "sshd"], check=True)

                self.app.notify("Root SSH login disabled. SSH reloaded.", severity="information")
                self.refresh_ssh_config()

            except Exception as e:
                self.app.notify(f"Failed to disable root login: {e}", severity="error")

        self.app.push_screen(ConfirmDialog(
            f"**Disable SSH login for root?**\n\n"
            f"This will prevent root from logging in via SSH.\n"
            f"You will need to use a regular user and sudo.\n\n"
            f"Sudo users with SSH key access:\n" + "\n".join(f"* {u}" for u in sudo_users_with_ssh),
            on_confirm=do_disable
        ))
